versionhq 1.1.11.6__tar.gz → 1.1.11.8__tar.gz

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/.github/workflows/run_tests.yml

@@ -8,13 +8,11 @@ permissions:
 env:
   LITELLM_API_KEY: ${{ secrets.LITELLM_API_KEY }}
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
-  OPENAI_API_BASE: ${{ secrets.OPENAI_API_BASE }}
   DEFAULT_REDIRECT_URL:  ${{ secrets.DEFAULT_REDIRECT_URL }}
   DEFAULT_USER_ID:  ${{ secrets.DEFAULT_USER_ID }}
   COMPOSIO_API_KEY: ${{ secrets.COMPOSIO_API_KEY }}
   DEFAULT_MODEL_NAME:  ${{ secrets.DEFAULT_MODEL_NAME }}
   GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
-  GEMINI_API_BASE: ${{ secrets.GEMINI_API_BASE }}
   ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
   ANTHROPIC_API_BASE: ${{ secrets.ANTHROPIC_API_BASE }}
   MEM0_API_KEY: ${{ secrets.MEM0_API_KEY }}

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/.gitignore

@@ -3,6 +3,8 @@ destinations.py
 
 entity_memory.py
 
+llm_connection_test.py
+
 train.py
 
 dist/

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: versionhq
-Version: 1.1.11.6
+Version: 1.1.11.8
 Summary: LLM orchestration frameworks for model-agnostic AI agents that handle complex outbound workflows
 Author-email: Kuriko Iwai <kuriko@versi0n.io>
 License: MIT License
@@ -319,8 +319,10 @@ src/
    pyenv install 3.12.8
    pyenv global 3.12.8  (optional: `pyenv global system` to get back to the system default ver.)
    uv python pin 3.12.8
+   echo 3.12.8 > .python-version
    ```
 
+
 3. Set up environment variables:
    Create a `.env` file in the project root and add the following:
    ```

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/README.md

@@ -243,8 +243,10 @@ src/
    pyenv install 3.12.8
    pyenv global 3.12.8  (optional: `pyenv global system` to get back to the system default ver.)
    uv python pin 3.12.8
+   echo 3.12.8 > .python-version
    ```
 
+
 3. Set up environment variables:
    Create a `.env` file in the project root and add the following:
    ```

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/pyproject.toml

@@ -15,7 +15,7 @@ exclude = ["test*", "__pycache__", "*.egg-info"]
 
 [project]
 name = "versionhq"
-version = "1.1.11.6"
+version = "1.1.11.8"
 authors = [{ name = "Kuriko Iwai", email = "kuriko@versi0n.io" }]
 description = "LLM orchestration frameworks for model-agnostic AI agents that handle complex outbound workflows"
 readme = "README.md"

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/src/versionhq/__init__.py

@@ -17,7 +17,7 @@ from versionhq.tool.model import Tool
 from versionhq.tool.composio_tool import ComposioHandler
 
 
-__version__ = "1.1.11.6"
+__version__ = "1.1.11.8"
 __all__ = [
     "Agent",
     "Customer",

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/src/versionhq/agent/model.py

@@ -6,10 +6,10 @@ from typing_extensions import Self
 from dotenv import load_dotenv
 import litellm
 
-from pydantic import UUID4, BaseModel, Field, InstanceOf, PrivateAttr, model_validator, field_validator, ConfigDict
+from pydantic import UUID4, BaseModel, Field, InstanceOf, PrivateAttr, model_validator, field_validator
 from pydantic_core import PydanticCustomError
 
-from versionhq.llm.model import LLM, DEFAULT_CONTEXT_WINDOW_SIZE, DEFAULT_MODEL_NAME
+from versionhq.llm.model import LLM, DEFAULT_CONTEXT_WINDOW_SIZE, DEFAULT_MODEL_NAME, PROVIDERS
 from versionhq.tool.model import Tool, ToolSet
 from versionhq.knowledge.model import BaseKnowledgeSource, Knowledge
 from versionhq.memory.contextual_memory import ContextualMemory
@@ -99,7 +99,7 @@ class Agent(BaseModel):
     tools: Optional[List[InstanceOf[Tool | ToolSet] | Type[Tool] | Any]] = Field(default_factory=list)
 
     # knowledge
-    knowledge_sources: Optional[List[BaseKnowledgeSource]] = Field(default=None)
+    knowledge_sources: Optional[List[BaseKnowledgeSource | Any]] = Field(default=None)
     _knowledge: Optional[Knowledge] = PrivateAttr(default=None)
 
     # memory
@@ -162,90 +162,44 @@ class Agent(BaseModel):
     @model_validator(mode="after")
     def set_up_llm(self) -> Self:
         """
-        Set up the base model and function calling model (if any) using the LLM class.
-        Pass the model config params: `llm`, `max_tokens`, `max_execution_time`, `callbacks`,`respect_context_window` to the LLM class.
-        The base model is selected on the client app, else use the default model.
+        Set up `llm` and `function_calling_llm` as valid LLM objects using the given values.
         """
-
         self.agent_ops_agent_name = self.role
+        self.llm = self._set_llm(llm=self.llm)
+        function_calling_llm = self.function_calling_llm if self.function_calling_llm else self.llm if self.llm else None
+        self.function_calling_llm = self._set_llm(llm=function_calling_llm)
+        return self
 
-        if isinstance(self.llm, LLM):
-            llm = self._set_llm_params(self.llm)
-            self.llm = llm
 
-        elif isinstance(self.llm, str) or self.llm is None:
-            model_name = self.llm if self.llm is not None else DEFAULT_MODEL_NAME
-            llm = LLM(model=model_name)
-            updated_llm = self._set_llm_params(llm)
-            self.llm = updated_llm
+    def _set_llm(self, llm: Any | None) -> LLM:
+        llm = llm if llm is not None else DEFAULT_MODEL_NAME
 
-        else:
-            if isinstance(self.llm, dict):
-                model_name = self.llm.pop("model_name", self.llm.pop("deployment_name", str(self.llm)))
-                llm = LLM(model=model_name if model_name is not None else DEFAULT_MODEL_NAME)
-                updated_llm = self._set_llm_params(llm, { k: v for k, v in self.llm.items() if v is not None })
-                self.llm = updated_llm
+        match llm:
+            case LLM():
+                return self._set_llm_params(llm=llm)
 
-            else:
+            case str():
+                llm_obj = LLM(model=llm)
+                return self._set_llm_params(llm=llm_obj)
+
+            case dict():
+                model_name = llm.pop("model_name", llm.pop("deployment_name", str(llm)))
+                llm_obj = LLM(model=model_name if model_name else DEFAULT_MODEL_NAME)
+                return self._set_llm_params(llm_obj, { k: v for k, v in llm.items() if v is not None })
+
+            case _:
                 model_name = (getattr(self.llm, "model_name") or getattr(self.llm, "deployment_name") or str(self.llm))
-                llm = LLM(model=model_name)
+                llm_obj = LLM(model=model_name)
                 llm_params = {
-                    "max_tokens": (getattr(self.llm, "max_tokens") or self.max_tokens or 3000),
-                    "timeout": getattr(self.llm, "timeout", self.max_execution_time),
-                    "callbacks": getattr(self.llm, "callbacks", None),
-                    "temperature": getattr(self.llm, "temperature", None),
-                    "logprobs": getattr(self.llm, "logprobs", None),
-                    "api_key": getattr(self.llm, "api_key", os.environ.get("LITELLM_API_KEY", None)),
-                    "base_url": getattr(self.llm, "base_url", None),
+                    "max_tokens": (getattr(llm, "max_tokens") or self.max_tokens or 3000),
+                    "timeout": getattr(llm, "timeout", self.max_execution_time),
+                    "callbacks": getattr(llm, "callbacks", None),
+                    "temperature": getattr(llm, "temperature", None),
+                    "logprobs": getattr(llm, "logprobs", None),
+                    "api_key": getattr(llm, "api_key", os.environ.get("LITELLM_API_KEY", None)),
+                    "base_url": getattr(llm, "base_url", None),
                 }
-                updated_llm = self._set_llm_params(llm, llm_params)
-                self.llm = updated_llm
-
-
-        """
-        Set up funcion_calling LLM as well.
-        Check if the model supports function calling, setup LLM instance accordingly, using the same params with the LLM.
-        """
-        if self.function_calling_llm:
-            if isinstance(self.function_calling_llm, LLM):
-                if self.function_calling_llm._supports_function_calling() == False:
-                    self.function_calling_llm = LLM(model=DEFAULT_MODEL_NAME)
-
-                updated_llm = self._set_llm_params(self.function_calling_llm)
-                self.function_calling_llm = updated_llm
-
-            elif isinstance(self.function_calling_llm, str):
-                llm = LLM(model=self.function_calling_llm)
-
-                if llm._supports_function_calling() == False:
-                    llm = LLM(model=DEFAULT_MODEL_NAME)
-
-                updated_llm = self._set_llm_params(llm)
-                self.function_calling_llm = updated_llm
-
-            else:
-                if isinstance(self.function_calling_llm, dict):
-                    model_name = self.function_calling_llm.pop("model_name", self.function_calling_llm.pop("deployment_name", str(self.function_calling_llm)))
-                    llm = LLM(model=model_name)
-                    updated_llm = self._set_llm_params(llm, { k: v for k, v in self.function_calling_llm.items() if v is not None })
-                    self.function_calling_llm = updated_llm
-
-                else:
-                    model_name = (getattr(self.function_calling_llm, "model_name") or getattr(self.function_calling_llm, "deployment_name") or str(self.function_calling_llm))
-                    llm = LLM(model=model_name)
-                    llm_params = {
-                        "max_tokens": (getattr(self.function_calling_llm, "max_tokens") or self.max_tokens or 3000),
-                        "timeout": getattr(self.function_calling_llm, "timeout", self.max_execution_time),
-                        "callbacks": getattr(self.function_calling_llm, "callbacks", None),
-                        "temperature": getattr(self.function_calling_llm, "temperature", None),
-                        "logprobs": getattr(self.function_calling_llm, "logprobs", None),
-                        "api_key": getattr(self.function_calling_llm, "api_key", os.environ.get("LITELLM_API_KEY", None)),
-                        "base_url": getattr(self.function_calling_llm, "base_url", None),
-                    }
-                    updated_llm = self._set_llm_params(llm, llm_params)
-                    self.function_calling_llm = updated_llm
-
-        return self
+                return self._set_llm_params(llm=llm_obj, config=llm_params)
 
 
     def _set_llm_params(self, llm: LLM, config: Dict[str, Any] = None) -> LLM:
@@ -257,6 +211,11 @@ class Agent(BaseModel):
         llm.timeout = self.max_execution_time if llm.timeout is None else llm.timeout
         llm.max_tokens = self.max_tokens if self.max_tokens else llm.max_tokens
 
+        if llm.provider is None:
+            provider_name = llm.model.split("/")[0]
+            valid_provider = provider_name if provider_name in PROVIDERS else None
+            llm.provider = valid_provider
+
         if self.callbacks:
             llm.callbacks = self.callbacks
             llm._set_callbacks(llm.callbacks)
@@ -344,14 +303,46 @@ class Agent(BaseModel):
 
     @model_validator(mode="after")
     def set_up_knowledge(self) -> Self:
-        if self.knowledge_sources:
-            collection_name = f"{self.role.replace(' ', '_')}"
+        from versionhq.knowledge.source import BaseKnowledgeSource, StringKnowledgeSource, TextFileKnowledgeSource, CSVKnowledgeSource, ExcelKnowledgeSource, JSONKnowledgeSource
+        from versionhq.knowledge.source_docling import DoclingSource
 
-            self._knowledge = Knowledge(
-                sources=self.knowledge_sources,
-                embedder_config=self.embedder_config,
-                collection_name=collection_name,
-            )
+        if self.knowledge_sources:
+            try:
+                collection_name = f"{self.role.replace(' ', '_')}"
+                knowledge_sources = []
+                docling_fp, txt_fp, json_fp, excel_fp, csv_fp, pdf_fp = [], [], [], [], [], []
+                str_cont = ""
+
+                for item in self.knowledge_sources:
+                    if isinstance(item, BaseKnowledgeSource):
+                        knowledge_sources.append(item)
+
+                    elif isinstance(item, str) and "http" in item:
+                        docling_fp.append(item)
+
+                    elif isinstance(item, str):
+                        match  os.path.splitext(item)[1]:
+                            case ".txt": txt_fp.append(item)
+                            case ".json": json_fp.append(item)
+                            case ".xls" | ".xlsx": excel_fp.append(item)
+                            case ".pdf": pdf_fp.append(item)
+                            case ".csv": csv_fp.append(item)
+                            case _: str_cont += str(item)
+
+                    else:
+                        str_cont += str(item)
+
+                if docling_fp: knowledge_sources.append(DoclingSource(file_paths=docling_fp))
+                if str_cont: knowledge_sources.append(StringKnowledgeSource(content=str_cont))
+                if txt_fp: knowledge_sources.append(TextFileKnowledgeSource(file_paths=txt_fp))
+                if csv_fp: knowledge_sources.append(CSVKnowledgeSource(file_path=csv_fp))
+                if excel_fp: knowledge_sources.append(ExcelKnowledgeSource(file_path=excel_fp))
+                if json_fp: knowledge_sources.append(JSONKnowledgeSource(file_paths=json_fp))
+
+                self._knowledge = Knowledge(sources=knowledge_sources, embedder_config=self.embedder_config, collection_name=collection_name)
+
+            except:
+                self._logger.log(level="warning", message="We cannot find the format for the source. Add BaseKnowledgeSource objects instead.", color="yellow")
 
         return self
 
@@ -414,7 +405,7 @@ class Agent(BaseModel):
             self._logger.log(level="info", message=f"Messages sent to the model: {messages}", color="blue")
 
             if tool_res_as_final:
-                func_llm = self.function_calling_llm if self.function_calling_llm and self.function_calling_llm._supports_function_calling() else LLM(model=DEFAULT_MODEL_NAME)
+                func_llm = self.function_calling_llm if self.function_calling_llm and self.function_calling_llm._supports_function_calling() else self.llm if self.llm and self.llm._supports_function_calling() else LLM(model=DEFAULT_MODEL_NAME)
                 raw_response = func_llm.call(messages=messages, tools=tools, tool_res_as_final=True)
                 task.tokens = func_llm._tokens
             else:
@@ -458,7 +449,7 @@ class Agent(BaseModel):
         from versionhq.knowledge._utils import extract_knowledge_context
 
         task: InstanceOf[Task] = task
-        tools: Optional[List[InstanceOf[Tool]| InstanceOf[ToolSet] | Type[Tool]]] = task_tools + self.tools if task.can_use_agent_tools else task_tools
+        tools: Optional[List[InstanceOf[Tool | ToolSet] | Type[Tool]]] = task_tools + self.tools if task.can_use_agent_tools else task_tools
 
         if self.max_rpm and self._rpm_controller:
             self._rpm_controller._reset_request_count()
@@ -474,7 +465,6 @@ class Agent(BaseModel):
                 if agent_knowledge_context:
                     task_prompt += agent_knowledge_context
 
-
         if self.use_memory == True:
             contextual_memory = ContextualMemory(
                 memory_config=self.memory_config, stm=self.short_term_memory, ltm=self.long_term_memory, um=self.user_memory

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/src/versionhq/knowledge/source_docling.py

@@ -3,7 +3,6 @@ from typing import Iterator, List, Optional
 from urllib.parse import urlparse
 
 try:
-    import docling
     from docling.datamodel.base_models import InputFormat
     from docling.document_converter import DocumentConverter
     from docling.exceptions import ConversionError
@@ -12,19 +11,12 @@ try:
     DOCLING_AVAILABLE = True
 except ImportError:
     import envoy
-    r = envoy.run("uv add docling --optional docling")
-
-    import docling
-    from docling.datamodel.base_models import InputFormat
-    from docling.document_converter import DocumentConverter
-    from docling.exceptions import ConversionError
-    from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker
-    from docling_core.types.doc.document import DoclingDocument
+    envoy.run("uv add docling --optional docling")
     DOCLING_AVAILABLE = True
 except:
     DOCLING_AVAILABLE = False
 
-from pydantic import Field, InstanceOf
+from pydantic import Field
 
 from versionhq.knowledge.source import BaseKnowledgeSource
 from versionhq.storage.utils import fetch_db_storage_path
@@ -54,11 +46,20 @@ class DoclingSource(BaseKnowledgeSource):
     ))
 
     def __init__(self, *args, **kwargs):
-        if not DOCLING_AVAILABLE:
-            raise ImportError("The docling package is required. Please install the package using: $ uv add docling.")
-        else:
+        if DOCLING_AVAILABLE:
+            from docling.datamodel.base_models import InputFormat
+            from docling.document_converter import DocumentConverter
+            from docling.exceptions import ConversionError
+            from docling_core.transforms.chunker.hierarchical_chunker import HierarchicalChunker
+            from docling_core.types.doc.document import DoclingDocument
+
             super().__init__(*args, **kwargs)
 
+        else:
+            raise ImportError("The docling package is required. Please install the package using: $ uv add docling.")
+        # else:
+        #     super().__init__(*args, **kwargs)
+
 
     def _convert_source_to_docling_documents(self) -> List["DoclingDocument"]:
         conv_results_iter = self.document_converter.convert_all(self.valid_file_paths)

{versionhq-1.1.11.6 → versionhq-1.1.11.8}/src/versionhq/llm/llm_vars.py

@@ -3,6 +3,33 @@ from typing import Type
 
 JSON_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
 
+PROVIDERS = [
+    "openai",
+    "gemini",
+    "sagemaker",
+    "huggingface", # need api base
+    "anthropic",
+    "ollama",
+    "watson",
+    "bedrock",
+    "azure",
+    "cerebras",
+    "llama",
+]
+
+ENDPOINT_PROVIDERS = [
+    # "openai",
+    # "gemini",
+    # "sagemaker",
+    "huggingface",
+    # "anthropic",
+    # "ollama",
+    # "watson",
+    # "bedrock",
+    # "azure",
+    # "cerebras",
+    # "llama",
+]
 
 """
 List of models available on the framework.
@@ -16,7 +43,6 @@ litellm.pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=
 
 MODELS = {
     "openai": [
-        # "gpt-3.5-turbo",
         "gpt-4",
         "gpt-4o",
         "gpt-4o-mini",
@@ -27,11 +53,7 @@ MODELS = {
         "gemini/gemini-1.5-flash",
         "gemini/gemini-1.5-pro",
         "gemini/gemini-2.0-flash-exp",
-        # "gemini/gemini-gemma-2-9b-it",
-        # "gemini/gemini-gemma-2-27b-it",
     ],
-    # "vetrex_ai": [
-    # ],
     "anthropic": [
         "claude-3-5-sonnet-20241022",
         "claude-3-5-sonnet-20240620",
@@ -39,10 +61,28 @@ MODELS = {
         "claude-3-opus-20240229",
         "claude-3-haiku-20240307",
     ],
-    # "ollama": [
-    #   "ollama/llama3.1",
-    #   "ollama/mixtral",
-    # ],
+    "huggingface": [
+        "huggingface/qwen/qwen2.5-VL-72B-Instruct",
+    ],
+    # "sagemaker": [
+    #     "sagemaker/huggingface-text2text-flan-t5-base",
+    #     "sagemaker/huggingface-llm-gemma-7b",
+    #     "sagemaker/jumpstart-dft-meta-textgeneration-llama-2-13b",
+    #     "sagemaker/jumpstart-dft-meta-textgeneration-llama-2-70b",
+    #     "sagemaker/jumpstart-dft-meta-textgeneration-llama-3-8b",
+    #     "sagemaker/jumpstart-dft-meta-textgeneration-llama-3-70b",
+    #     "sagemaker/huggingface-llm-mistral-7b"
+    # ], #https://docs.aws.amazon.com/sagemaker/latest/dg/jumpstart-foundation-models-latest.html
+    "ollama": [
+      "ollama/llama3.1",
+      "ollama/mixtral",
+      "ollama/mixtral-8x22B-Instruct-v0.1",
+    ],
+    "deepseek": [
+        "deepseek/deepseek-reasoner",
+
+    ],
+
     # "watson": [
     #     "watsonx/meta-llama/llama-3-1-70b-instruct",
     #     "watsonx/meta-llama/llama-3-1-8b-instruct",
@@ -53,44 +93,48 @@ MODELS = {
     #     "watsonx/mistral/mistral-large",
     #     "watsonx/ibm/granite-3-8b-instruct",
     # ],
-    # "bedrock": [
-    #     "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
-    #     "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
-    #     "bedrock/anthropic.claude-3-haiku-20240307-v1:0",
-    #     "bedrock/anthropic.claude-3-opus-20240229-v1:0",
-    #     "bedrock/anthropic.claude-v2:1",
-    #     "bedrock/anthropic.claude-v2",
-    #     "bedrock/anthropic.claude-instant-v1",
-    #     "bedrock/meta.llama3-1-405b-instruct-v1:0",
-    #     "bedrock/meta.llama3-1-70b-instruct-v1:0",
-    #     "bedrock/meta.llama3-1-8b-instruct-v1:0",
-    #     "bedrock/meta.llama3-70b-instruct-v1:0",
-    #     "bedrock/meta.llama3-8b-instruct-v1:0",
-    #     "bedrock/amazon.titan-text-lite-v1",
-    #     "bedrock/amazon.titan-text-express-v1",
-    #     "bedrock/cohere.command-text-v14",
-    #     "bedrock/ai21.j2-mid-v1",
-    #     "bedrock/ai21.j2-ultra-v1",
-    #     "bedrock/ai21.jamba-instruct-v1:0",
-    #     "bedrock/meta.llama2-13b-chat-v1",
-    #     "bedrock/meta.llama2-70b-chat-v1",
-    #     "bedrock/mistral.mistral-7b-instruct-v0:2",
-    #     "bedrock/mistral.mixtral-8x7b-instruct-v0:1",
-    # ],
+    "bedrock": [
+        "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
+        "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+        "bedrock/anthropic.claude-3-haiku-20240307-v1:0",
+        "bedrock/anthropic.claude-3-opus-20240229-v1:0",
+        # "bedrock/anthropic.claude-v2:1",
+        "bedrock/anthropic.claude-v2",
+        "bedrock/anthropic.claude-instant-v1",
+        "bedrock/meta.llama3-1-405b-instruct-v1:0",
+        "bedrock/meta.llama3-1-70b-instruct-v1:0",
+        "bedrock/meta.llama3-1-8b-instruct-v1:0",
+        "bedrock/meta.llama3-70b-instruct-v1:0",
+        "bedrock/meta.llama3-8b-instruct-v1:0",
+        "bedrock/amazon.titan-text-lite-v1",
+        "bedrock/amazon.titan-text-express-v1",
+        "bedrock/cohere.command-text-v14",
+        "bedrock/ai21.j2-mid-v1",
+        "bedrock/ai21.j2-ultra-v1",
+        "bedrock/ai21.jamba-instruct-v1:0",
+        "bedrock/meta.llama2-13b-chat-v1",
+        "bedrock/meta.llama2-70b-chat-v1",
+        "bedrock/mistral.mistral-7b-instruct-v0:2",
+        "bedrock/mistral.mixtral-8x7b-instruct-v0:1",
+    ],
 }
 
 
-PROVIDERS = [
-    "openai",
-    "anthropic",
-    "gemini",
-    "ollama",
-    "watson",
-    "bedrock",
-    "azure",
-    "cerebras",
-    "llama",
-]
+
+KEYS = {
+    "openai": ["OPENAI_API_KEY"],
+    "gemini": ["GEMINI_API_KEY"],
+    "sagemaker": ["AWS_ACCESS_KEY_ID", "ADW_SECURET_ACCESS_KEY", "AWS_REGION_NAME"],
+    "anthropic": ["ANTHROPIC_API_KEY"],
+}
+
+
+"""
+Use base_url to specify
+"""
+BASE_URLS = {
+    "deepseek": "https://api.deepseek.com"
+}
 
 
 """
@@ -118,6 +162,8 @@ LLM_CONTEXT_WINDOW_SIZES = {
     "claude-3-haiku-20240307": 200000,
 
     "deepseek-chat": 128000,
+    "deepseek/deepseek-reasoner": 8192,
+
     "gemma2-9b-it": 8192,
     "gemma-7b-it": 8192,
     "llama3-groq-70b-8192-tool-use-preview": 8192,
@@ -135,11 +181,7 @@ LLM_CONTEXT_WINDOW_SIZES = {
 }
 
 
-LLM_API_KEY_NAMES = {
-    "openai":  "OPENAI_API_KEY",
-    "anthropic": "ANTHROPIC_API_KEY",
-    "gemini": "GEMINI_API_KEY",
-}
+
 
 LLM_BASE_URL_KEY_NAMES = {
     "openai":  "OPENAI_API_BASE",
@@ -262,14 +304,8 @@ PARAMS = {
     ],
     "openai": [
         "timeout",
-        # "temperature",
-        # "top_p",
-        # "n",
-        # "stream",
         "stream_options",
-        # "stop",
         "max_compl,etion_tokens",
-        # "max_tokens",
         "modalities",
         "prediction",
         "audio",
@@ -277,10 +313,7 @@ PARAMS = {
         "frequency_penalty",
         "logit_bias",
         "user",
-        # "response_format",
         "seed",
-        # "tools",
-        # "tool_choice",
         "logprobs",
         "top_logprobs",
         "parallel_tool_calls",