versionhq 1.1.10.8__py3-none-any.whl → 1.1.10.9__py3-none-any.whl

versionhq/knowledge/storage.py

@@ -0,0 +1,177 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+import contextlib
+import hashlib
+import io
+import logging
+import os
+import shutil
+
+import chromadb
+import chromadb.errors
+from chromadb.api import ClientAPI
+from chromadb.api.types import OneOrMany
+from chromadb.config import Settings
+
+from versionhq._utils.logger import Logger
+from versionhq._utils.vars import KNOWLEDGE_DIRECTORY
+from versionhq.storage.utils import fetch_db_storage_path
+from versionhq.knowledge.embedding import EmbeddingConfigurator
+
+
+@contextlib.contextmanager
+def suppress_logging(
+    logger_name="chromadb.segment.impl.vector.local_persistent_hnsw",
+    level=logging.ERROR,
+):
+    logger = logging.getLogger(logger_name)
+    original_level = logger.getEffectiveLevel()
+    logger.setLevel(level)
+    with (
+        contextlib.redirect_stdout(io.StringIO()),
+        contextlib.redirect_stderr(io.StringIO()),
+        contextlib.suppress(UserWarning),
+    ):
+        yield
+    logger.setLevel(original_level)
+
+
+
+class BaseKnowledgeStorage(ABC):
+    """
+    Abstract base class for knowledge storage implementations.
+    """
+
+    @abstractmethod
+    def search(self, query: List[str], limit: int = 3, filter: Optional[dict] = None, score_threshold: float = 0.35) -> List[Dict[str, Any]]:
+        """Search for documents in the knowledge base."""
+        pass
+
+    @abstractmethod
+    def save(self, documents: List[str], metadata: Dict[str, Any] | List[Dict[str, Any]]) -> None:
+        """Save documents to the knowledge base."""
+        pass
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the knowledge base."""
+        pass
+
+
+
+class KnowledgeStorage(BaseKnowledgeStorage):
+    """
+    Extends Storage to handle embeddings for memory entries, improving search efficiency.
+    """
+
+    collection: Optional[chromadb.Collection] = None
+    collection_name: Optional[str] = "knowledge"
+    app: Optional[ClientAPI] = None
+
+    def __init__(self, embedder_config: Optional[Dict[str, Any]] = None, collection_name: Optional[str] = None):
+        self.collection_name = collection_name
+        self._set_embedder_config(embedder_config)
+
+    def search(self, query: List[str], limit: int = 3, filter: Optional[dict] = None, score_threshold: float = 0.35) -> List[Dict[str, Any]]:
+        with suppress_logging():
+            if self.collection:
+                fetched = self.collection.query(query_texts=query, n_results=limit, where=filter)
+                results = []
+                for i in range(len(fetched["ids"][0])):
+                    result = {
+                        "id": fetched["ids"][0][i],
+                        "metadata": fetched["metadatas"][0][i],
+                        "context": fetched["documents"][0][i],
+                        "score": fetched["distances"][0][i],
+                    }
+                    if result["score"] >= score_threshold:
+                        results.append(result)
+                return results
+            else:
+                raise Exception("Collection not initialized")
+
+
+    def initialize_knowledge_storage(self):
+        base_path = os.path.join(fetch_db_storage_path(), "knowledge")
+        chroma_client = chromadb.PersistentClient(path=base_path, settings=Settings(allow_reset=True))
+        self.app = chroma_client
+
+        try:
+            collection_name = f"knowledge_{self.collection_name}" if self.collection_name else "knowledge"
+            if self.app:
+                self.collection = self.app.get_or_create_collection(name=collection_name, embedding_function=self.embedder_config)
+            else:
+                raise Exception("Vector Database Client not initialized")
+        except Exception:
+            raise Exception("Failed to create or get collection")
+
+
+    def reset(self):
+        base_path = os.path.join(fetch_db_storage_path(), KNOWLEDGE_DIRECTORY)
+        if not self.app:
+            self.app = chromadb.PersistentClient(path=base_path, settings=Settings(allow_reset=True))
+        self.app.reset()
+        shutil.rmtree(base_path)
+        self.app = None
+        self.collection = None
+
+
+    def save(self, documents: List[str], metadata: Optional[Dict[str, Any] | List[Dict[str, Any]]] = None) -> None:
+        if not self.collection:
+            raise Exception("Collection not initialized")
+
+        try:
+            unique_docs = {}
+            for i, doc in enumerate(documents):
+                doc_id = hashlib.sha256(doc.encode("utf-8")).hexdigest()
+                doc_metadata = None
+                if metadata is not None:
+                    if isinstance(metadata, list):
+                        doc_metadata = metadata[i]
+                    else:
+                        doc_metadata = metadata
+                unique_docs[doc_id] = (doc, doc_metadata)
+
+            filtered_docs = []
+            filtered_metadata = []
+            filtered_ids = []
+
+            for doc_id, (doc, meta) in unique_docs.items():
+                filtered_docs.append(doc)
+                filtered_metadata.append(meta)
+                filtered_ids.append(doc_id)
+
+            final_metadata: Optional[OneOrMany[chromadb.Metadata]] = (
+                None if all(m is None for m in filtered_metadata) else filtered_metadata
+            )
+            self.collection.upsert(documents=filtered_docs, metadatas=final_metadata, ids=filtered_ids)
+
+        except chromadb.errors.InvalidDimensionException as e:
+            Logger(verbose=True).log(
+                level="error",
+                message="Embedding dimension mismatch. This usually happens when mixing different embedding models.",
+                color="red",
+            )
+            raise ValueError("Embedding dimension mismatch. Make sure you're using the same embedding model across all operations with this collection.") from e
+
+        except Exception as e:
+            Logger(verbose=True).log(level="error", message=f"Failed to upsert documents: {str(e)}", color="red")
+            raise
+
+
+    def _create_default_embedding_function(self) -> Any:
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
+        )
+
+
+    def _set_embedder_config(self, embedder_config: Optional[Dict[str, Any]] = None) -> None:
+        """
+        Set the embedding configuration for the knowledge storage.
+        """
+        self.embedder_config = EmbeddingConfigurator().configure_embedder(embedder_config) if embedder_config else self._create_default_embedding_function()

versionhq/llm/model.py

@@ -82,6 +82,7 @@ class LLM(BaseModel):
 
     _logger: Logger = PrivateAttr(default_factory=lambda: Logger(verbose=True))
     _init_model_name: str = PrivateAttr(default=None)
+    _tokens: int = PrivateAttr(default=0) # accumulate total tokens used for the call
     model_config = ConfigDict(extra="allow")
 
     model: str = Field(default=DEFAULT_MODEL_NAME)
@@ -136,6 +137,7 @@ class LLM(BaseModel):
 
         self._init_model_name = self.model
         self.model = None
+        self._tokens = 0
 
         if self.provider and MODELS.get(self.provider):
             provider_model_list = MODELS.get(self.provider)
@@ -147,7 +149,7 @@ class LLM(BaseModel):
                         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._logger.log(level="warning", 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:
@@ -163,7 +165,7 @@ class LLM(BaseModel):
                             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._logger.log(level="warning", 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"
 
@@ -185,17 +187,36 @@ class LLM(BaseModel):
         return self
 
 
+    def _create_valid_params(self, config: Dict[str, Any], provider: str = None) -> Dict[str, Any]:
+        params = dict()
+        valid_keys = list()
+
+        if not provider:
+            valid_keys = PARAMS.get("litellm") + PARAMS.get("common") + PARAMS.get(self.provider) if self.provider else PARAMS.get("litellm") + PARAMS.get("common")
+        else:
+            valid_keys = PARAMS.get("common") + PARAMS.get(self.provider)
+
+        for item in valid_keys:
+            if hasattr(self, item) and getattr(self, item):
+                params[item] = getattr(self, item)
+            elif item in config:
+                params[item] = config[item]
+
+        return params
+
+
     def call(
         self,
         messages: List[Dict[str, str]],
         response_format: Optional[Dict[str, Any]] = None,
-        tools: Optional[List[Tool | ToolSet | Type[Tool]]] = None,
+        tools: Optional[List[Tool | ToolSet | Any ]] = 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 the agent's params and model params.
         """
+
         litellm.drop_params = True
 
         with suppress_warnings():
@@ -203,80 +224,73 @@ class LLM(BaseModel):
                 self._set_callbacks(self.callbacks) # passed by agent
 
             try:
-                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 tool_res_as_final else response_format
-
                 provider = self.provider if self.provider else "openai"
+                self.response_format = { "type": "json_object" } if tool_res_as_final == True else response_format
 
-                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:
-                    messages.append(res["choices"][0]["message"])
-                    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))))
+                if not tools:
+                    params = self._create_valid_params(config=config)
+                    res = litellm.completion(messages=messages, stream=False, **params)
+                    self._tokens += int(res["usage"]["total_tokens"])
+                    return res["choices"][0]["message"]["content"]
 
-                        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)
+                else:
+                    self.tools = [item.tool.properties if isinstance(item, ToolSet) else item.properties for item in tools]
 
-                                if tool_res_as_final:
-                                    tool_res += str(res)
-                                else:
-                                    messages.append({ "role": "tool", "tool_call_id": item.id, "content": str(res) })
+                    if provider == "openai":
+                        params = self._create_valid_params(config=config, provider=provider)
+                        res = openai_client.chat.completions.create(messages=messages, model=self.model, tools=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):
+                                try:
+                                    func_args = json.loads(json.dumps(eval(str(func_args))))
+                                except:
+                                    pass
+
+                            for tool in tools:
+                                if isinstance(tool, ToolSet) and (tool.tool.name == func_name or tool.tool.func.__name__ == func_name or func_name == "random_func"):
+                                    tool_instance = tool.tool
+                                    args = tool.kwargs
+                                    tool_res_to_add = tool_instance.run(params=args)
+
+                                    if tool_res_as_final:
+                                        tool_res += str(tool_res_to_add)
+                                    else:
+                                        messages.append(res.choices[0].message)
+                                        messages.append({ "role": "tool", "tool_call_id": item.id, "content": str(tool_res_to_add) })
 
-                            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)
-                                if tool_res_as_final:
-                                    tool_res += str(res)
                                 else:
-                                    messages.append({ "role": "tool", "tool_call_id": item.id, "content": str(res) })
-
-                    if tool_res_as_final:
-                        return tool_res
-
-                    else:
-                        res = litellm.completion(messages=messages, stream=False, **params)
-
-                return res["choices"][0]["message"]["content"]
-
+                                    try:
+                                        tool_res_to_add = tool.run(params=func_args)
+                                        if tool_res_as_final:
+                                            tool_res += str(tool_res_to_add)
+                                        else:
+                                            messages.append(res.choices[0].message)
+                                            messages.append({ "role": "tool", "tool_call_id": item.id, "content": str(tool_res_to_add) })
+                                    except:
+                                        pass
+
+                        if tool_res_as_final:
+                            return tool_res
+                        else:
+                            res = openai_client.chat.completions.create(messages=messages, model=self.model, tools=self.tools)
+                            self._tokens += int(res["usage"]["total_tokens"])
+                            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
+                raise e
 
             except Exception as 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:
         try:

versionhq/memory/contextual_memory.py

@@ -0,0 +1,96 @@
+from typing import Any, Dict, Optional, List
+
+from versionhq.memory.model import ShortTermMemory, LongTermMemory, UserMemory
+
+
+class ContextualMemory:
+    """
+    A class to construct context from memories (ShortTermMemory, UserMemory).
+    The context will be added to the prompt when the agent executes the task.
+    """
+
+    def __init__(
+        self,
+        memory_config: Optional[Dict[str, Any]],
+        stm: ShortTermMemory,
+        ltm: LongTermMemory,
+        um: UserMemory,
+        # em: EntityMemory,
+    ):
+        self.memory_provider = memory_config.get("provider") if memory_config is not None else None
+        self.stm = stm
+        self.ltm = ltm
+        self.um = um
+
+
+    def build_context_for_task(self, task, context: List[Any] | str) -> str:
+        """
+        Automatically builds a minimal, highly relevant set of contextual information for a given task.
+        """
+
+        query = f"{task.description} {context}".strip()
+
+        if query == "":
+            return ""
+
+        context = []
+        context.append(self._fetch_stm_context(query))
+        if self.memory_provider == "mem0":
+            context.append(self._fetch_user_context(query))
+        return "\n".join(filter(None, context))
+
+
+    def _fetch_stm_context(self, query) -> str:
+        """
+        Fetches recent relevant insights from STM related to the task's description and expected_output, formatted as bullet points.
+        """
+        stm_results = self.stm.search(query)
+        formatted_results = "\n".join(
+            [
+                f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
+                for result in stm_results
+            ]
+        )
+        return f"Recent Insights:\n{formatted_results}" if stm_results else ""
+
+
+    def _fetch_ltm_context(self, task) -> Optional[str]:
+        """
+        Fetches historical data or insights from LTM that are relevant to the task's description and expected_output, formatted as bullet points.
+        """
+        ltm_results = self.ltm.search(task, latest_n=2)
+        if not ltm_results:
+            return None
+
+        formatted_results = [suggestion for result in ltm_results for suggestion in result["metadata"]["suggestions"]]
+        formatted_results = list(dict.fromkeys(formatted_results))
+        formatted_results = "\n".join([f"- {result}" for result in formatted_results])
+        return f"Historical Data:\n{formatted_results}" if ltm_results else ""
+
+
+    def _fetch_user_context(self, query: str) -> str:
+        """
+        Fetches and formats relevant user information from User Memory.
+        """
+
+        user_memories = self.um.search(query)
+        if not user_memories:
+            return ""
+
+        formatted_memories = "\n".join(f"- {result['memory']}" for result in user_memories)
+        return f"User memories/preferences:\n{formatted_memories}"
+
+
+    # def _fetch_entity_context(self, query) -> str:
+    #     """
+    #     Fetches relevant entity information from Entity Memory related to the task's description and expected_output,
+    #     formatted as bullet points.
+    #     """
+    #     em_results = self.em.search(query)
+    #     formatted_results = "\n".join(
+    #         [
+    #             f"- {result['memory'] if self.memory_provider == 'mem0' else result['context']}"
+    #             for result in em_results
+    #         ]  # type: ignore #  Invalid index type "str" for "str"; expected type "SupportsIndex | slice"
+    #     )
+    #     return f"Entities:\n{formatted_results}" if em_results else ""

versionhq/memory/model.py

@@ -0,0 +1,174 @@
+from typing import Any, Dict, List, Optional
+
+from versionhq.storage.rag_storage import RAGStorage
+from versionhq.storage.ltm_sqlite_storage import LTMSQLiteStorage
+
+
+class Memory:
+    """
+    Base class for memory.
+    """
+
+    def __init__(self, storage: RAGStorage):
+        self.storage = storage
+
+
+    def save(self, value: Any, metadata: Optional[Dict[str, Any]] = None, agent: Optional[str] = None) -> None:
+        metadata = metadata or {}
+
+        if agent:
+            metadata["agent"] = agent
+        self.storage.save(value, metadata)
+
+
+    def search(self, query: str, limit: int = 3, score_threshold: float = 0.35) -> List[Any]:
+        return self.storage.search(query=query, limit=limit, score_threshold=score_threshold)
+
+
+
+class ShortTermMemoryItem:
+    def __init__(
+        self,
+        data: Any,
+        agent: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        self.data = data
+        self.agent = agent
+        self.metadata = metadata if metadata is not None else {}
+
+
+class ShortTermMemory(Memory):
+    """
+    A class for managing transient data related to immediate tasks and interactions.
+    - Type: stm
+    - Storage: Mem0Storage | RAGStorage
+    """
+
+    def __init__(self, agent = None, embedder_config: Dict[str, Any] = None, storage=None, path=None):
+        if hasattr(agent, "memory_config") and agent.memory_config is not None:
+            self.memory_provider = agent.memory_config.get("provider")
+        else:
+            self.memory_provider = None
+
+        if self.memory_provider == "mem0":
+            try:
+                from versionhq.storage.mem0_storage import Mem0Storage
+            except ImportError:
+                raise ImportError("Mem0 is not installed. Please install it with `uv pip install mem0ai`.")
+
+            storage = Mem0Storage(type="stm", agent=agent)
+        else:
+            storage = storage if storage else RAGStorage(type="stm", embedder_config=embedder_config, agents=[agent,], path=path)
+
+        super().__init__(storage)
+
+
+    def save(self, value: Any, metadata: Optional[Dict[str, Any]] = None, agent: Optional[str] = None) -> None:
+        item = ShortTermMemoryItem(data=value, metadata=metadata, agent=agent)
+        if self.memory_provider == "mem0":
+            item.data = f"Remember the following insights from Agent run: {item.data}"
+
+        super().save(value=item.data, metadata=item.metadata, agent=item.agent)
+
+
+    def search(self, query: str, limit: int = 3, score_threshold: float = 0.35,):
+        return self.storage.search(query=query, limit=limit, score_threshold=score_threshold)
+
+
+    def reset(self) -> None:
+        try:
+            self.storage.reset()
+        except Exception as e:
+            raise Exception(f"An error occurred while resetting the short-term memory: {str(e)}")
+
+
+
+class LongTermMemoryItem:
+    def __init__(
+        self,
+        agent: str,
+        task: str,
+        datetime: str,
+        quality: Optional[int | float] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        self.task = task
+        self.agent = agent
+        self.quality = quality
+        self.datetime = datetime
+        self.metadata = metadata if metadata is not None else {}
+
+
+
+class LongTermMemory(Memory):
+    """
+    A class for managing cross runs data related to overall task executions.
+    - Type: ltm
+    - Storage: LTMSQLiteStorage
+    """
+
+    def __init__(self, storage=None, path=None):
+        if not storage:
+            storage = LTMSQLiteStorage(db_path=path) if path else LTMSQLiteStorage()
+
+        super().__init__(storage)
+
+
+    def save(self, item: LongTermMemoryItem) -> None:
+        metadata = item.metadata
+        metadata.update({ "agent": item.agent })
+        self.storage.save(
+            task_description=item.task,
+            score=metadata["quality"],
+            metadata=metadata,
+            datetime=item.datetime,
+        )
+
+
+    def search(self, task: str, latest_n: int = 3) -> List[Dict[str, Any]]:
+        return self.storage.load(task, latest_n)
+
+
+    def reset(self) -> None:
+        self.storage.reset()
+
+
+
+class UserMemoryItem:
+    def __init__(self, data: Any, user: str, metadata: Optional[Dict[str, Any]] = None):
+        self.data = data
+        self.user = user
+        self.metadata = metadata if metadata is not None else {} # can be stored last purchased item, comm related to the user
+
+
+class UserMemory(Memory):
+    """
+    UserMemory class for handling user memory storage and retrieval.
+    - Type: user
+    - Storage: Mem0Storage
+    - Requirements: `user_id` in metadata
+    """
+
+    def __init__(self, agent=None, user_id=None):
+        try:
+            from versionhq.storage.mem0_storage import Mem0Storage
+        except ImportError:
+            raise ImportError("Mem0 is not installed. Please install it with `uv pip install mem0ai`.")
+
+        if not user_id:
+            raise ValueError("Need User Id to create UserMemory.")
+
+        else:
+            storage = Mem0Storage(type="user", agent=agent, user_id=user_id)
+            super().__init__(storage)
+
+
+    def save(self, value: str, metadata: Optional[Dict[str, Any]] = None, agent: Optional[str] = None) -> None:
+        data = f"Remember the details about the user: {value}"
+        super().save(value=data, metadata=metadata, agent=agent)
+
+
+    def search(self, query: str, limit: int = 3, score_threshold: float = 0.35):
+        results = self.storage.search(query=query, limit=limit, score_threshold=score_threshold)
+        return results

versionhq/storage/base.py

@@ -0,0 +1,14 @@
+from typing import Any, Dict, List
+
+
+class Storage:
+    """Abstract base class defining the storage interface"""
+
+    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+        pass
+
+    def search(self, query: str, limit: int, score_threshold: float) -> Dict[str, Any] | List[Any]:
+        return {}
+
+    def reset(self) -> None:
+        pass