Move constants to constants.py, slight refactor

Author: natn0

letta/constants.py

@@ -10,6 +10,7 @@
 
 ADMIN_PREFIX = "/v1/admin"
 API_PREFIX = "/v1"
+OLLAMA_API_PREFIX = "/v1"
 OPENAI_API_PREFIX = "/openai"
 
 COMPOSIO_ENTITY_ENV_VAR_KEY = "COMPOSIO_ENTITY"
@@ -50,8 +51,9 @@
 # Max steps for agent loop
 DEFAULT_MAX_STEPS = 50
 
-# minimum context window size
+# context window size
 MIN_CONTEXT_WINDOW = 4096
+DEFAULT_CONTEXT_WINDOW = 32000
 
 # number of concurrent embedding requests to sent
 EMBEDDING_BATCH_SIZE = 200
@@ -63,6 +65,7 @@
 # embeddings
 MAX_EMBEDDING_DIM = 4096  # maximum supported embeding size - do NOT change or else DBs will need to be reset
 DEFAULT_EMBEDDING_CHUNK_SIZE = 300
+DEFAULT_EMBEDDING_DIM = 1024
 
 # tokenizers
 EMBEDDING_TO_TOKENIZER_MAP = {

letta/schemas/providers/ollama.py

@@ -3,7 +3,7 @@
 import aiohttp
 from pydantic import Field
 
-from letta.constants import DEFAULT_EMBEDDING_CHUNK_SIZE
+from letta.constants import DEFAULT_EMBEDDING_CHUNK_SIZE, DEFAULT_CONTEXT_WINDOW, DEFAULT_EMBEDDING_DIM, OLLAMA_API_PREFIX
 from letta.log import get_logger
 from letta.schemas.embedding_config import EmbeddingConfig
 from letta.schemas.enums import ProviderCategory, ProviderType
@@ -12,8 +12,6 @@
 
 logger = get_logger(__name__)
 
-ollama_prefix = "/v1"
-
 
 class OllamaProvider(OpenAIProvider):
     """Ollama provider that uses the native /api/generate endpoint
@@ -41,8 +39,9 @@ async def list_llm_models_async(self) -> list[LLMConfig]:
                 response_json = await response.json()
 
         configs = []
-        for model in response_json["models"]:
-            model_details = await self._get_model_details_async(model["name"])
+        for model in response_json.get("models", []):
+            model_name = model["name"]
+            model_details = await self._get_model_details_async(model_name)
             if not model_details or "completion" not in model_details.get("capabilities", []):
                 continue
 
@@ -53,17 +52,17 @@ async def list_llm_models_async(self) -> list[LLMConfig]:
                     context_window = int(context_length)
 
             if context_window is None:
-                print(f"Ollama model {model['name']} has no context window, using default 32000")
-                context_window = 32000
+                logger.warning(f"Ollama model {model_name} has no context window, using default {DEFAULT_CONTEXT_WINDOW}")
+                context_window = DEFAULT_CONTEXT_WINDOW
 
             configs.append(
                 LLMConfig(
-                    model=model["name"],
+                    model=model_name,
                     model_endpoint_type=ProviderType.ollama,
-                    model_endpoint=f"{self.base_url}{ollama_prefix}",
+                    model_endpoint=f"{self.base_url}{OLLAMA_API_PREFIX}",
                     model_wrapper=self.default_prompt_formatter,
                     context_window=context_window,
-                    handle=self.get_handle(model["name"]),
+                    handle=self.get_handle(model_name),
                     provider_name=self.name,
                     provider_category=self.provider_category,
                 )
@@ -83,8 +82,9 @@ async def list_embedding_models_async(self) -> list[EmbeddingConfig]:
                 response_json = await response.json()
 
         configs = []
-        for model in response_json["models"]:
-            model_details = await self._get_model_details_async(model["name"])
+        for model in response_json.get("models", []):
+            model_name = model["name"]
+            model_details = await self._get_model_details_async(model_name)
             if not model_details or "embedding" not in model_details.get("capabilities", []):
                 continue
 
@@ -95,17 +95,17 @@ async def list_embedding_models_async(self) -> list[EmbeddingConfig]:
                     embedding_dim = int(embedding_length)
 
             if not embedding_dim:
-                print(f"Ollama model {model['name']} has no embedding dimension, using default 1024")
-                embedding_dim = 1024
+                logger.warning(f"Ollama model {model_name} has no embedding dimension, using default {DEFAULT_EMBEDDING_DIM}")
+                embedding_dim = DEFAULT_EMBEDDING_DIM
 
             configs.append(
                 EmbeddingConfig(
-                    embedding_model=model["name"],
+                    embedding_model=model_name,
                     embedding_endpoint_type=ProviderType.ollama,
-                    embedding_endpoint=f"{self.base_url}{ollama_prefix}",
+                    embedding_endpoint=f"{self.base_url}{OLLAMA_API_PREFIX}",
                     embedding_dim=embedding_dim,
                     embedding_chunk_size=DEFAULT_EMBEDDING_CHUNK_SIZE,
-                    handle=self.get_handle(model["name"], is_embedding=True),
+                    handle=self.get_handle(model_name, is_embedding=True),
                 )
             )
         return configs
@@ -125,5 +125,4 @@ async def _get_model_details_async(self, model_name: str) -> dict | None:
                     return await response.json()
         except Exception as e:
             logger.warning(f"Failed to get model details for {model_name} with error: {e}")
-
-        return None
+            return None