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

versionhq/storage/ltm_sqlite_storage.py

@@ -0,0 +1,131 @@
+import json
+import sqlite3
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from versionhq._utils.logger import Logger
+from versionhq.storage.utils import fetch_db_storage_path
+
+
+class LTMSQLiteStorage:
+    """
+    An updated SQLite storage class for LTM data storage.
+    """
+
+    def __init__(self, db_path: Optional[str] = None) -> None:
+        if db_path is None:
+            db_path = str(Path(fetch_db_storage_path()) / "ltm_storage.db")
+
+        self.db_path = db_path
+        self._logger: Logger = Logger(verbose=True)
+
+        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
+        self._initialize_db()
+
+
+    def _initialize_db(self):
+        """
+        Initializes the SQLite database and creates LTM table
+        """
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS long_term_memories (
+                        id INTEGER PRIMARY KEY AUTOINCREMENT,
+                        task_description TEXT,
+                        metadata TEXT,
+                        datetime TEXT,
+                        score REAL
+                    )
+                """
+                )
+
+                conn.commit()
+
+        except sqlite3.Error as e:
+            self._logger.log(
+                level="error",
+                message=f"MEMORY ERROR: An error occurred during database initialization: {str(e)}",
+                color="red",
+            )
+
+    def save(self, task_description: str, metadata: Dict[str, Any], datetime: str, score: int | float) -> None:
+        """
+        Saves data to the LTM table with error handling.
+        """
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                INSERT INTO long_term_memories (task_description, metadata, datetime, score)
+                VALUES (?, ?, ?, ?)
+            """,
+                    (task_description, json.dumps(metadata), datetime, score),
+                )
+                conn.commit()
+        except sqlite3.Error as e:
+            self._logger.log(
+                level="error",
+                message=f"MEMORY ERROR: An error occurred while saving to LTM: {str(e)}",
+                color="red",
+            )
+
+
+    def load(self, task_description: str, latest_n: int) -> Optional[List[Dict[str, Any]]]:
+        """
+        Queries the LTM table by task description with error handling.
+        """
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    f"""
+                    SELECT metadata, datetime, score
+                    FROM long_term_memories
+                    WHERE task_description = ?
+                    ORDER BY datetime DESC, score ASC
+                    LIMIT {latest_n}
+                """,
+                    (task_description,),
+                )
+                rows = cursor.fetchall()
+                if rows:
+                    return [
+                        {
+                            "metadata": json.loads(row[0]),
+                            "datetime": row[1],
+                            "score": row[2],
+                        }
+                        for row in rows
+                    ]
+
+        except sqlite3.Error as e:
+            self._logger.log(
+                level="error",
+                message=f"MEMORY ERROR: An error occurred while querying LTM: {e}",
+                color="red",
+            )
+        return None
+
+
+    def reset(self) -> None:
+        """
+        Resets the LTM table with error handling.
+        """
+
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute("DELETE FROM long_term_memories")
+                conn.commit()
+
+        except sqlite3.Error as e:
+            self._logger.log(
+                level="error",
+                message=f"MEMORY ERROR: An error occurred while deleting all rows in LTM: {str(e)}",
+                color="red",
+            )
+        return None

versionhq/storage/mem0_storage.py

@@ -0,0 +1,109 @@
+import os
+from typing import Any, Dict, List
+
+from mem0 import MemoryClient
+
+from versionhq.storage.base import Storage
+
+
+class Mem0Storage(Storage):
+    """
+    Extends Storage to handle embedding and searching across entities using Mem0.
+    """
+
+    def __init__(self, type, agent=None, user_id=None):
+        """
+        Create a memory client using API keys and other config.
+        """
+
+        super().__init__()
+
+        if type not in ["user", "stm", "ltm", "entities"]:
+            raise ValueError("Invalid type for Mem0Storage. Must be 'user' or 'agent'.")
+
+        self.memory_type = type
+        self.agent= agent
+        self.memory_config = agent.memory_config
+
+        user_id = user_id if user_id else self._get_user_id()
+        if type == "user" and not user_id:
+            raise ValueError("User ID is required for user memory type")
+
+        config = self.memory_config.get("config", {})
+        mem0_api_key = os.environ.get("MEM0_API_KEY", config.get("api_key"))
+        mem0_org_id = config.get("org_id")
+        mem0_project_id = config.get("project_id")
+
+        if mem0_org_id and mem0_project_id:
+            self.memory = MemoryClient(api_key=mem0_api_key, org_id=mem0_org_id, project_id=mem0_project_id)
+        else:
+            self.memory = MemoryClient(api_key=mem0_api_key)
+
+
+    def _sanitize_role(self, role: str) -> str:
+        """
+        Sanitizes agent roles to ensure valid directory names.
+        """
+        return role.replace("\n", "").replace(" ", "_").replace("/", "_")
+
+
+    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+        user_id = self._get_user_id()
+        agent_name = self._get_agent_name()
+
+        if self.memory_type == "user":
+            self.memory.add(value, user_id=user_id, metadata={**metadata})
+
+        elif self.memory_type == "stm":
+            agent_name = self._get_agent_name()
+            self.memory.add(value, agent_id=agent_name, metadata={"type": "stm", **metadata})
+
+        elif self.memory_type == "ltm":
+            agent_name = self._get_agent_name()
+            self.memory.add(value, agent_id=agent_name, infer=False, metadata={"type": "ltm", **metadata})
+
+        elif self.memory_type == "entities":
+            entity_name = self._get_agent_name()
+            self.memory.add(value, user_id=entity_name, metadata={"type": "entity", **metadata})
+
+
+    def search(self, query: str, limit: int = 3, score_threshold: float = 0.35) -> List[Any]:
+        params = {"query": query, "limit": limit}
+
+        if self.memory_type == "user":
+            user_id = self._get_user_id()
+            params["user_id"] = user_id
+
+        elif self.memory_type == "stm":
+            agent_name = self._get_agent_name()
+            params["agent_id"] = agent_name
+            params["metadata"] = {"type": "stm"}
+
+        elif self.memory_type == "ltm":
+            agent_name = self._get_agent_name()
+            params["agent_id"] = agent_name
+            params["metadata"] = {"type": "ltm"}
+
+        elif self.memory_type == "entities":
+            agent_name = self._get_agent_name()
+            params["agent_id"] = agent_name
+            params["metadata"] = {"type": "entity"}
+
+        results = self.memory.search(**params)
+        return [r for r in results if r["score"] >= score_threshold]
+
+
+    def _get_user_id(self):
+        if self.memory_type == "user":
+            if hasattr(self, "memory_config") and self.memory_config is not None:
+                return self.memory_config.get("config", {}).get("user_id")
+            else:
+                return None
+        return None
+
+
+    def _get_agent_name(self):
+        agents = self.agents if self.agents else []
+        agents = [self._sanitize_role(agent.role) for agent in agents]
+        agents = "_".join(agents)
+        return agents

versionhq/storage/rag_storage.py

@@ -0,0 +1,231 @@
+import contextlib
+import io
+import logging
+import os
+import shutil
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+from chromadb.api import ClientAPI
+
+from versionhq.knowledge.embedding import EmbeddingConfigurator
+from versionhq._utils.vars import MAX_FILE_NAME_LENGTH
+from versionhq.storage.utils import fetch_db_storage_path
+
+
+@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 BaseRAGStorage(ABC):
+    """
+    Base class for RAG-based Storage implementations.
+    """
+
+    app: Any | None = None
+
+    def __init__(
+        self,
+        type: str,
+        allow_reset: bool = True,
+        embedder_config: Optional[Any] = None,
+        agents: List[Any] = None,
+    ):
+        self.type = type
+        self.allow_reset = allow_reset
+        self.embedder_config = embedder_config
+        self.agents = agents
+
+    def _initialize_agents(self) -> str:
+        if self.agents:
+            return "_".join(
+                [self._sanitize_role(agent.role) for agent in self.agents]
+            )
+        return ""
+
+    @abstractmethod
+    def _sanitize_role(self, role: str) -> str:
+        """Sanitizes agent roles to ensure valid directory names."""
+        pass
+
+    @abstractmethod
+    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+        """Save a value with metadata to the storage."""
+        pass
+
+    @abstractmethod
+    def search(self, query: str, limit: int = 3, filter: Optional[dict] = None, score_threshold: float = 0.35) -> List[Any]:
+        """Search for entries in the storage."""
+        pass
+
+    @abstractmethod
+    def reset(self) -> None:
+        """Reset the storage."""
+        pass
+
+    @abstractmethod
+    def _generate_embedding(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> Any:
+        """Generate an embedding for the given text and metadata."""
+        pass
+
+    @abstractmethod
+    def _initialize_app(self):
+        """Initialize the vector db."""
+        pass
+
+    def setup_config(self, config: Dict[str, Any]):
+        """Setup the config of the storage."""
+        pass
+
+    def initialize_client(self):
+        """Initialize the client of the storage. This should setup the app and the db collection"""
+        pass
+
+
+
+class RAGStorage(BaseRAGStorage):
+    """
+    Extends Storage to handle embeddings for memory entries, improving
+    search efficiency.
+    """
+
+    app: ClientAPI | None = None
+
+    def __init__(self, type, allow_reset=True, embedder_config=None, agents=list(), path=None):
+        super().__init__(type, allow_reset, embedder_config, agents)
+        agents = agents
+        agents = [self._sanitize_role(agent.role) for agent in agents]
+        agents = "_".join(agents)
+
+        self.agents = agents
+        self.storage_file_name = self._build_storage_file_name(type, agents)
+        self.type = type
+        self.allow_reset = allow_reset
+        self.path = path
+        self._initialize_app()
+
+
+    def _set_embedder_config(self):
+        configurator = EmbeddingConfigurator()
+        self.embedder_config = configurator.configure_embedder(self.embedder_config)
+
+
+    def _initialize_app(self) -> None:
+        import chromadb
+        from chromadb.config import Settings
+
+        self._set_embedder_config()
+        chroma_client = chromadb.PersistentClient(
+            path=self.path if self.path else self.storage_file_name,
+            settings=Settings(allow_reset=self.allow_reset),
+        )
+        self.app = chroma_client
+
+        try:
+            self.collection = self.app.get_collection(name=self.type, embedding_function=self.embedder_config)
+        except Exception:
+            self.collection = self.app.create_collection(name=self.type, embedding_function=self.embedder_config)
+
+
+    def _sanitize_role(self, role: str) -> str:
+        """
+        Sanitizes agent roles to ensure valid directory names.
+        """
+        return role.replace("\n", "").replace(" ", "_").replace("/", "_")
+
+
+    def _build_storage_file_name(self, type: str, file_name: str) -> str:
+        """
+        Ensures file name does not exceed max allowed by OS
+        """
+        base_path = f"{fetch_db_storage_path()}/{type}"
+
+        if len(file_name) > MAX_FILE_NAME_LENGTH:
+            logging.warning(
+                f"Trimming file name from {len(file_name)} to {MAX_FILE_NAME_LENGTH} characters."
+            )
+            file_name = file_name[:MAX_FILE_NAME_LENGTH]
+
+        return f"{base_path}/{file_name}"
+
+
+    def save(self, value: Any, metadata: Dict[str, Any]) -> None:
+        if not hasattr(self, "app") or not hasattr(self, "collection"):
+            self._initialize_app()
+        try:
+            self._generate_embedding(value, metadata)
+        except Exception as e:
+            logging.error(f"Error during {self.type} save: {str(e)}")
+
+
+    def search(self, query: str, limit: int = 3, filter: Optional[dict] = None, score_threshold: float = 0.35) -> List[Any]:
+        if not hasattr(self, "app"):
+            self._initialize_app()
+
+        try:
+            with suppress_logging():
+                response = self.collection.query(query_texts=query, n_results=limit)
+
+            results = []
+            for i in range(len(response["ids"][0])):
+                result = {
+                    "id": response["ids"][0][i],
+                    "metadata": response["metadatas"][0][i],
+                    "context": response["documents"][0][i],
+                    "score": response["distances"][0][i],
+                }
+                if result["score"] >= score_threshold:
+                    results.append(result)
+
+            return results
+        except Exception as e:
+            logging.error(f"Error during {self.type} search: {str(e)}")
+            return []
+
+
+    def _generate_embedding(self, text: str, metadata: Dict[str, Any]) -> None:
+        if not hasattr(self, "app") or not hasattr(self, "collection"):
+            self._initialize_app()
+
+        self.collection.add(
+            documents=[text],
+            metadatas=[metadata or {}],
+            ids=[str(uuid.uuid4())],
+        )
+
+
+    def reset(self) -> None:
+        try:
+            if self.app:
+                self.app.reset()
+                shutil.rmtree(f"{fetch_db_storage_path()}/{self.type}")
+                self.app = None
+                self.collection = None
+        except Exception as e:
+            if "attempt to write a readonly database" in str(e):
+                pass
+            else:
+                raise Exception(f"An error occurred while resetting the {self.type} memory: {e}")
+
+    def _create_default_embedding_function(self):
+        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"
+        )

versionhq/storage/task_output_storage.py

@@ -1,21 +1,10 @@
-import appdirs
-import os
 import json
 import sqlite3
 import datetime
 from typing import Any, Dict, List, Optional
-from dotenv import load_dotenv
-from pathlib import Path
 
 from versionhq._utils.logger import Logger
-
-load_dotenv(override=True)
-
-def fetch_db_storage_path():
-    directory_name = Path.cwd().name
-    data_dir = Path(appdirs.user_data_dir(appname=directory_name, appauthor="Version IO Sdn Bhd.", version=None, roaming=False))
-    data_dir.mkdir(parents=True, exist_ok=True)
-    return data_dir
+from versionhq.storage.utils import fetch_db_storage_path
 
 storage_path = fetch_db_storage_path()
 default_db_name = "task_outputs"
@@ -37,25 +26,25 @@ class TaskOutputSQLiteStorage:
         Initializes the SQLite database and creates LTM table.
         """
 
-        # try:
-        with sqlite3.connect(self.db_path) as conn:
-            cursor = conn.cursor()
-            cursor.execute(
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS task_outputs (
+                        task_id TEXT PRIMARY KEY,
+                        output JSON,
+                        task_index INTEGER,
+                        inputs JSON,
+                        was_replayed BOOLEAN,
+                        timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
+                    )
                 """
-                CREATE TABLE IF NOT EXISTS task_outputs (
-                    task_id TEXT PRIMARY KEY,
-                    output JSON,
-                    task_index INTEGER,
-                    inputs JSON,
-                    was_replayed BOOLEAN,
-                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
                 )
-            """
-            )
-            conn.commit()
+                conn.commit()
 
-        # except sqlite3.Error as e:
-        #     self._logger.log(level="error", message=f"DATABASE INITIALIZATION ERROR: {e}", color="red")
+        except sqlite3.Error as e:
+            self._logger.log(level="error", message=f"SQL database initialization failed: {str(e)}", color="red")
 
 
     def add(self, task, output: Dict[str, Any], task_index: int, was_replayed: bool = False, inputs: Dict[str, Any] = {}):
@@ -91,7 +80,7 @@ class TaskOutputSQLiteStorage:
 
                 if cursor.rowcount == 0:
                     self._logger.log(
-                        level="info", message=f"No row found with task_index {task_index}. No update performed.", color="yellow",
+                        level="warning", message=f"No row found with task_index {task_index}. No update performed.", color="yellow",
                     )
 
         except sqlite3.Error as e:

versionhq/storage/utils.py

@@ -0,0 +1,26 @@
+import appdirs
+import os
+from dotenv import load_dotenv
+from pathlib import Path
+
+load_dotenv(override=True)
+
+def fetch_db_storage_path() -> str:
+    directory_name = get_project_directory_name()
+    data_dir = Path(appdirs.user_data_dir(appname=directory_name, appauthor="Version IO Sdn Bhd", version=None, roaming=False))
+    data_dir.mkdir(parents=True, exist_ok=True)
+    return str(data_dir)
+
+
+def get_project_directory_name() -> str:
+    """
+    Returns the current project directory name
+    """
+    project_directory_name = os.environ.get("STORAGE_DIR")
+
+    if project_directory_name:
+        return project_directory_name
+    else:
+        cwd = Path.cwd()
+        project_directory_name = cwd.name
+        return project_directory_name

versionhq/task/TEMPLATES/Description.py

@@ -0,0 +1,5 @@
+EVALUATE="""Evaluate the accuracy and quality of the following task output. For each criterion defined below, provide a score between 0 (completely inaccurate/unacceptable) and 1 (perfectly accurate/acceptable), and suggest specific improvements to address any shortcomings.
+Task: {task_description}
+Task Output: {task_output}
+Evaluation criteria: {eval_criteria}
+"""

versionhq/task/evaluate.py

@@ -0,0 +1,122 @@
+from typing import List, Optional, Dict, Any
+from typing_extensions import Self
+
+from pydantic import BaseModel, Field, InstanceOf, model_validator
+
+"""
+Evaluate task output from accuracy, token consumption, latency perspectives, and mark the score from 0 to 1.
+"""
+
+
+
+
+class ScoreFormat:
+    def __init__(self, rate: float | int = 0, weight: int = 1):
+        self.rate = rate
+        self.weight = weight
+        self.aggregate = rate * weight
+
+
+class Score:
+    """
+    Evaluate the score on 0 (no performance) to 1 scale.
+    `rate`: Any float from 0.0 to 1.0 given by an agent.
+    `weight`: Importance of each factor to the aggregated score.
+    """
+
+    def __init__(
+        self,
+        brand_tone: ScoreFormat = ScoreFormat(0, 0),
+        audience: ScoreFormat = ScoreFormat(0, 0),
+        track_record: ScoreFormat = ScoreFormat(0, 0),
+        config: Optional[Dict[str, ScoreFormat]] = None
+    ):
+        self.brand_tone = brand_tone
+        self.audience = audience
+        self.track_record = track_record
+        self.config = config
+
+        if self.config:
+            for k, v in self.config.items():
+                if isinstance(v, ScoreFormat):
+                    setattr(self, k, v)
+
+
+    def result(self) -> int:
+        aggregate_score, denominator = 0, 0
+
+        for k, v in self.__dict__.items():
+            aggregate_score += v.aggregate
+            denominator += v.weight
+
+        if denominator == 0:
+            return 0
+
+        return round(aggregate_score / denominator, 2)
+
+
+class EvaluationItem(BaseModel):
+    """
+    A class to store evaluation and suggestion by the given criteria such as accuracy.
+    """
+    criteria: str
+    suggestion: str
+    score: int | float
+
+    def _convert_score_to_score_format(self, weight: int = 1) -> ScoreFormat | None:
+        if self.score and isinstance(self.score, (int, float)):
+            return ScoreFormat(rate=self.score, weight=weight)
+
+        else: return None
+
+
+
+class Evaluation(BaseModel):
+    # expected_outcome: Optional[str] = Field(default=None, description="human input on expected outcome")
+    items: List[EvaluationItem] = []
+    latency: int = Field(default=None, description="seconds")
+    tokens: int = Field(default=None, description="tokens consumed")
+    responsible_agent: Any = Field(default=None, description="store agent instance that evaluates the outcome")
+
+    @model_validator(mode="after")
+    def set_up_responsible_agent(self) -> Self:
+        from versionhq.agent.default_agents import task_evaluator
+        self.responsible_agent = task_evaluator
+        return self
+
+
+    @property
+    def aggregate_score(self) -> float:
+        """
+        Calcurate aggregate score from evaluation items.
+        """
+        if not self.items:
+            return 0
+
+        aggregate_score = 0
+        denominator = 0
+
+        for item in self.items:
+            score_format = item._convert_score_to_score_format()
+            aggregate_score += score_format.aggregate if score_format else 0
+            denominator += score_format.weight if score_format else 0
+
+        if denominator == 0:
+            return 0
+
+        return round(aggregate_score / denominator, 2)
+
+
+    @property
+    def suggestion_summary(self) -> str | None:
+        """
+        Return a summary of the suggestions
+        """
+        if not self.items:
+            return None
+
+        summary = ""
+        for item in self.items:
+            summary += f"{item.suggestion}, "
+
+        return summary