versionhq 1.1.9.13__py3-none-any.whl → 1.1.10.2__py3-none-any.whl

versionhq/llm/model.py

@@ -1,22 +1,38 @@
 import logging
+import json
 import os
 import sys
 import threading
 import warnings
 import litellm
+from litellm import JSONSchemaValidationError
+from abc import ABC
 from dotenv import load_dotenv
 from litellm import get_supported_openai_params
 from contextlib import contextmanager
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Type
+from typing_extensions import Self
 
-from versionhq.llm.llm_vars import LLM_CONTEXT_WINDOW_SIZES
+from pydantic import UUID4, BaseModel, Field, PrivateAttr, field_validator, model_validator, create_model, InstanceOf, ConfigDict
+from pydantic_core import PydanticCustomError
+
+from openai import OpenAI
+
+from versionhq.llm.llm_vars import LLM_CONTEXT_WINDOW_SIZES, LLM_API_KEY_NAMES, LLM_BASE_URL_KEY_NAMES, MODELS, PARAMS, SchemaType
 from versionhq.task import TaskOutputFormat
-from versionhq.task.model import ResponseField
+from versionhq.task.model import ResponseField, Task
+from versionhq.tool.model import Tool, ToolSet
+from versionhq._utils.logger import Logger
+
 
 load_dotenv(override=True)
-API_KEY_LITELLM = os.environ.get("API_KEY_LITELLM")
-DEFAULT_CONTEXT_WINDOW = int(8192 * 0.75)
-os.environ["LITELLM_LOG"] = "DEBUG"
+LITELLM_API_KEY = os.environ.get("LITELLM_API_KEY")
+LITELLM_API_BASE = os.environ.get("LITELLM_API_BASE")
+DEFAULT_CONTEXT_WINDOW_SIZE = int(8192 * 0.75)
+DEFAULT_MODEL_NAME = os.environ.get("DEFAULT_MODEL_NAME")
+
+proxy_openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"), organization="versionhq", base_url=LITELLM_API_BASE)
+openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
 
 class FilteredStream:
@@ -56,184 +72,231 @@ def suppress_warnings():
             sys.stderr = old_stderr
 
 
-class LLMResponseSchema:
-    """
-    Use the response schema for LLM response.
-    `field_list` contains the title, value type, bool if required of each field that needs to be returned.
-    field_list: [{ title, type, required } ]
-
-    i.e., reponse_schema
-    response_type: "array"  *options: "array", "dict"
-    propeties: { "recipe_name": { "type": "string" }, },
-    required:  ["recipe_name"]
-    """
-
-    def __init__(self, response_type: str, field_list: List[ResponseField]):
-        self.type = response_type
-        self.field_list = field_list
-
-    @property
-    def schema(self):
-        if len(self.field_list) == 0:
-            return
-
-        properties = [
-            {
-                field.title: {
-                    "type": field.type,
-                }
-            }
-            for field in self.field_list
-        ]
-        required = [field.title for field in self.field_list if field.required == True]
-        response_schema = {
-            "type": self.type,
-            "items": {"type": "object", "properties": {*properties}},
-            "required": required,
-        }
-        return response_schema
-
-
-class LLM:
+class LLM(BaseModel):
     """
+    An LLM class to store params except for response formats which will be given in the task handling process.
     Use LiteLLM to connect with the model of choice.
-    (Memo) Response formats will be given at the Task handling.
+    Some optional params are passed by the agent, else follow the default settings of the model provider.
+    Ref. https://docs.litellm.ai/docs/completion/input
     """
 
-    def __init__(
-        self,
-        model: str,
-        timeout: Optional[float | int] = None,
-        max_tokens: Optional[int] = None,
-        max_completion_tokens: Optional[int] = None,
-        context_window_size: Optional[int] = DEFAULT_CONTEXT_WINDOW,
-        callbacks: List[Any] = [],
-        temperature: Optional[float] = None,
-        top_p: Optional[float] = None,
-        n: Optional[int] = None,
-        stop: Optional[str | List[str]] = None,
-        presence_penalty: Optional[float] = None,
-        frequency_penalty: Optional[float] = None,
-        logit_bias: Optional[Dict[int, float]] = None,
-        # response_format: Optional[Dict[str, Any]] = None,
-        seed: Optional[int] = None,
-        logprobs: Optional[bool] = None,
-        top_logprobs: Optional[int] = None,
-        base_url: Optional[str] = None,
-        api_version: Optional[str] = None,
-        api_key: Optional[str] = None,
-        **kwargs,
-    ):
-        self.model = model
-        self.timeout = timeout
-        self.max_tokens = max_tokens
-        self.max_completion_tokens = max_completion_tokens
-        self.context_window_size = context_window_size
-        self.callbacks = callbacks
-
-        self.temperature = temperature
-        self.top_p = top_p
-        self.n = n
-        self.stop = stop
-        self.presence_penalty = presence_penalty
-        self.frequency_penalty = frequency_penalty
-        self.logit_bias = logit_bias
-        # self.response_format = response_format
-        self.seed = seed
-        self.logprobs = logprobs
-        self.top_logprobs = top_logprobs
-
-        self.base_url = base_url
-        self.api_version = api_version
-        self.api_key = api_key if api_key else API_KEY_LITELLM
-
-        self.kwargs = kwargs
-
-        litellm.drop_params = True
-        litellm.set_verbose=True
-        self.set_callbacks(callbacks)
+    _logger: Logger = PrivateAttr(default_factory=lambda: Logger(verbose=True))
+    _init_model_name: str = PrivateAttr(default=None)
+    model_config = ConfigDict(extra="allow")
+
+    model: str = Field(default=DEFAULT_MODEL_NAME)
+    provider: Optional[str] = Field(default=None, description="model provider or custom model provider")
+    base_url: Optional[str] = Field(default=None, description="api base of the model provider")
+    api_key: Optional[str] = Field(default=None, description="api key of the model provider")
+
+    # optional params
+    timeout: Optional[float | int] = Field(default=None)
+    max_tokens: Optional[int] = Field(default=None)
+    max_completion_tokens: Optional[int] = Field(default=None)
+    context_window_size: Optional[int] = Field(default=DEFAULT_CONTEXT_WINDOW_SIZE)
+    callbacks: List[Any] = Field(default_factory=list)
+    temperature: Optional[float] = Field(default=None)
+    top_p: Optional[float] = Field(default=None)
+    n: Optional[int] = Field(default=None)
+    stop: Optional[str | List[str]] = Field(default=None)
+    presence_penalty: Optional[float] = Field(default=None)
+    frequency_penalty: Optional[float] = Field(default=None)
+    logit_bias: Optional[Dict[int, float]] = Field(default=None)
+    seed: Optional[int] = Field(default=None)
+    logprobs: Optional[bool] = Field(default=None)
+    top_logprobs: Optional[int] = Field(default=None)
+    response_format: Optional[Any] = Field(default=None)
+    tools: Optional[List[Dict[str, Any]]] = Field(default_factory=list, description="store a list of tool properties")
+
+    # LiteLLM specific fields
+    api_base: Optional[str] = Field(default=None, description="litellm specific field - api base of the model provider")
+    api_version: Optional[str] = Field(default=None)
+    num_retries: Optional[int] = Field(default=2)
+    context_window_fallback_dict: Optional[Dict[str, Any]] = Field(default=None, description="A mapping of model to use if call fails due to context window error")
+    fallbacks: Optional[List[Any]]= Field(default=None, description="A list of model names + params to be used, in case the initial call fails")
+    metadata: Optional[Dict[str, Any]] = Field(default=None)
+
+    litellm.drop_params = True
+    litellm.set_verbose = True
+    os.environ['LITELLM_LOG'] = 'DEBUG'
+
+    @model_validator(mode="after")
+    def validate_base_params(self) -> Self:
+        """
+        1) Set up a valid model name with the provider name using the MODEL list.
+        * Assign a default model and provider based on the given information when no model key is found in the MODEL list.
+
+        2) Set up other base parameters for the model and LiteLLM.
+        """
+
+        if self.model is None:
+            self._logger.log(level="error", message="Model name is missing.", color="red")
+            raise PydanticCustomError("model_missing", "The model name must be provided.", {})
+
+
+        self._init_model_name = self.model
+        self.model = None
+
+        if self.provider and MODELS.get(self.provider):
+            provider_model_list = MODELS.get(self.provider)
+            for item in provider_model_list:
+                if self.model is None:
+                    if item == self._init_model_name:
+                        self.model = item
+                    elif self._init_model_name in item and self.model is None:
+                        self.model = item
+                    else:
+                        temp_model = provider_model_list[0]
+                        self._logger.log(level="info", message=f"The provided model: {self._init_model_name} is not in the list. We'll assign a model: {temp_model} from the selected model provider: {self.provider}.", color="yellow")
+                        self.model = temp_model
+
+        else:
+            for k, v in MODELS.items():
+                for item in v:
+                    if self.model is None:
+                        if self._init_model_name == item:
+                            self.model = item
+                            self.provider = k
+
+                        elif self.model is None and self._init_model_name in item:
+                            self.model = item
+                            self.provider = k
+
+            if self.model is None:
+                self._logger.log(level="info", message=f"The provided model \'{self.model}\' is not in the list. We'll assign a default model.", color="yellow")
+                self.model = DEFAULT_MODEL_NAME
+                self.provider = "openai"
+
+
+        if self.callbacks:
+            self._set_callbacks(self.callbacks)
+
+        self.context_window_size = self._get_context_window_size()
+
+        api_key_name = self.provider.upper() + "_API_KEY" if self.provider else None
+        if api_key_name:
+            self.api_key = os.environ.get(api_key_name, None)
+
+        base_url_key_name = self.provider.upper() + "_API_BASE" if self.provider else None
+        if base_url_key_name:
+            self.base_url = os.environ.get(base_url_key_name)
+            self.api_base = self.base_url
+
+        return self
+
 
     def call(
         self,
-        output_formats: List[TaskOutputFormat],
-        field_list: Optional[List[ResponseField]],
         messages: List[Dict[str, str]],
-        callbacks: List[Any] = [],
+        response_format: Optional[Dict[str, Any]] = None,
+        tools: Optional[List[Tool | ToolSet | Type[Tool]]] = None,
+        config: Optional[Dict[str, Any]] = {}, # any other conditions to pass on to the model.
+        tool_res_as_final: bool = False
     ) -> str:
         """
-        Execute LLM based on Agent's controls.
+        Execute LLM based on the agent's params and model params.
         """
 
         with suppress_warnings():
-            if callbacks and len(callbacks) > 0:
-                self.set_callbacks(callbacks)
+            if len(self.callbacks) > 0:
+                self._set_callbacks(self.callbacks)
 
             try:
-                response_format = None
-
-                #! REFINEME
-                if TaskOutputFormat.JSON in output_formats:
-                    response_format = LLMResponseSchema(
-                        response_type="json_object", field_list=field_list
-                    )
-
-                params = {
-                    "model": self.model,
-                    "messages": messages,
-                    "timeout": self.timeout,
-                    "temperature": self.temperature,
-                    "top_p": self.top_p,
-                    "n": self.n,
-                    "stop": self.stop,
-                    "max_tokens": self.max_tokens or self.max_completion_tokens,
-                    "presence_penalty": self.presence_penalty,
-                    "frequency_penalty": self.frequency_penalty,
-                    "logit_bias": self.logit_bias,
-                    # "response_format": response_format,
-                    "seed": self.seed,
-                    "logprobs": self.logprobs,
-                    "top_logprobs": self.top_logprobs,
-                    "api_base": self.base_url,
-                    "api_version": self.api_version,
-                    "api_key": self.api_key,
-                    "stream": False,
-                    **self.kwargs,
-                }
-                params = {k: v for k, v in params.items() if v is not None}
-                res = litellm.completion(**params)
+                if tools:
+                    self.tools = [item.tool.properties if isinstance(item, ToolSet) else item.properties for item in tools]
+
+                if response_format:
+                    self.response_format = { "type": "json_object" } if self.model == "gpt-3.5-turbo" or tool_res_as_final else response_format
+
+                provider = self.provider if self.provider else "openai"
+
+                params = {}
+                valid_params = PARAMS.get("litellm") + PARAMS.get("common") + PARAMS.get(self.provider) if self.provider else PARAMS.get("litellm") + PARAMS.get("common")
+
+                for item in valid_params:
+                    if item:
+                        if hasattr(self, item) and getattr(self, item):
+                            params[item] = getattr(self, item)
+                        elif item in config:
+                            params[item] = config[item]
+                        else:
+                            continue
+                    else:
+                        continue
+
+                res = litellm.completion(messages=messages, stream=False, **params)
+
+                if self.tools:
+                    tool_calls = res["choices"][0]["message"]["tool_calls"]
+                    tool_res = ""
+
+                    for item in tool_calls:
+                        func_name = item.function.name
+                        func_args = item.function.arguments
+
+                        if not isinstance(func_args, dict):
+                            func_args = json.loads(json.dumps(eval(str(func_args))))
+
+                        for tool in tools:
+                            if isinstance(tool, ToolSet) and (tool.tool.name.replace(" ", "_") == func_name or tool.tool.func.__name__ == func_name):
+                                tool_instance = tool.tool
+                                args = tool.kwargs
+                                res = tool_instance.run(params=args)
+                                tool_res += str(res)
+
+                            elif (isinstance(tool, Tool) or type(tool) == Tool) and (tool.name.replace(" ", "_") == func_name or tool.func.__name__ == func_name):
+                                res = tool.run(params=func_args)
+                                tool_res += str(res)
+
+                    if tool_res_as_final == True:
+                        return tool_res
+                        pass
+
+                    else:
+                        messages.append({ "role": "tool", "tool_call_id": tool_calls.id, "content": tool_res })
+                        res = litellm.completion(messages=messages, stream=False, **params)
+
                 return res["choices"][0]["message"]["content"]
 
+
+            except JSONSchemaValidationError as e:
+                self._logger.log(level="error", message="Raw Response: {}".format(e.raw_response), color="red")
+                return None
+
             except Exception as e:
-                logging.error(f"LiteLLM call failed: {str(e)}")
+                self._logger.log(level="error", message=f"{self.model} failed to execute: {str(e)}", color="red")
+                if "litellm.RateLimitError" in str(e):
+                    raise e
+
                 return None
 
-    def supports_function_calling(self) -> bool:
+
+    def _supports_function_calling(self) -> bool:
         try:
             params = get_supported_openai_params(model=self.model)
             return "response_format" in params
         except Exception as e:
-            logging.error(f"Failed to get supported params: {str(e)}")
+            self._logger.log(level="error", message=f"Failed to get supported params: {str(e)}", color="red")
             return False
 
-    def supports_stop_words(self) -> bool:
+
+    def _supports_stop_words(self) -> bool:
         try:
             params = get_supported_openai_params(model=self.model)
             return "stop" in params
         except Exception as e:
-            logging.error(f"Failed to get supported params: {str(e)}")
+            self._logger.log(level="error", message=f"Failed to get supported params: {str(e)}", color="red")
             return False
 
-    def get_context_window_size(self) -> int:
+
+    def _get_context_window_size(self) -> int:
         """
         Only use 75% of the context window size to avoid cutting the message in the middle.
         """
-        return (
-            int(LLM_CONTEXT_WINDOW_SIZES.get(self.model) * 0.75)
-            if hasattr(LLM_CONTEXT_WINDOW_SIZES, self.model)
-            else DEFAULT_CONTEXT_WINDOW
-        )
+        return int(LLM_CONTEXT_WINDOW_SIZES.get(self.model) * 0.75) if LLM_CONTEXT_WINDOW_SIZES.get(self.model) is not None else DEFAULT_CONTEXT_WINDOW_SIZE
+
 
-    def set_callbacks(self, callbacks: List[Any]):
+    def _set_callbacks(self, callbacks: List[Any]):
         callback_types = [type(callback) for callback in callbacks]
         for callback in litellm.success_callback[:]:
             if type(callback) in callback_types:

versionhq/task/__init__.py

@@ -5,5 +5,5 @@ class TaskOutputFormat(str, Enum):
     """Enum that represents the output format of a task."""
 
     JSON = "json"
-    PYDANTIC = "pydantic"
+    PYDANTIC = "pydantic model"
     RAW = "raw"