Merge pull request #4 from XortexAI/feat/classifier

Migration: Phase 4.1, Part of Phase 5.2, Part of Phase 6, Part of Phase 8, Phase 9.3

Author: ved015

.gitignore

@@ -1 +1,47 @@
-.env
+# Python
+*.py[cod]
+__pycache__/
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual environments
+venv/
+.venv/
+env/
+.env
+
+# Pytest
+.pytest_cache/
+.coverage
+htmlcov/
+coverage.xml
+
+# IDEs
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# MacOS
+.DS_Store
+
+# Logs
+logs/
+*.log

pyproject.toml

@@ -7,7 +7,7 @@ name = "Xmem"
 version = "0.1.0"
 description = "Universal unified memory system for AI agents"
 authors = [
-    {name = "Vedant Mahajan", email = "xmemlabs@gmail.com"}
+    {name = "Vedant Mahajan", email = "xmemlabs@gmail.com"},
     {name = "Ishaan Gupta", email = "xmemlabs@gmail.com"}
 ]
 readme = "README.md"

src/agents/__init__.py

@@ -0,0 +1,7 @@
+"""Xmem agents — re-export the public agent classes."""
+
+from src.agents.classifier import ClassifierAgent
+
+__all__ = [
+    "ClassifierAgent",
+]

src/agents/base.py

@@ -0,0 +1,35 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from langchain_core.language_models import BaseChatModel
+
+
+class BaseAgent(ABC):
+    def __init__(self, model: BaseChatModel, name: str, system_prompt: str = ""):
+        self.model = model
+        self.name = name
+        self.system_prompt = system_prompt
+        self.logger = logging.getLogger(f"xmem.agents.{name}")
+
+    @abstractmethod
+    async def arun(self, state: Dict[str, Any]) -> Any:
+        ...
+
+    def run(self, state: Dict[str, Any]) -> Any:
+        import asyncio
+        return asyncio.run(self.arun(state))
+
+    def _build_messages(self, user_message: str) -> list:
+        messages = []
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+        messages.append({"role": "user", "content": user_message})
+        return messages
+
+    async def _call_model(self, messages: list) -> str:
+        response = await self.model.ainvoke(messages)
+        content = response.content
+        if isinstance(content, list):
+            content = "\n".join(str(c) for c in content)
+        return content

src/agents/classifier.py

@@ -0,0 +1,51 @@
+"""
+Classifier Agent — the entry-point router for Xmem.
+
+Classifies user input into one or more intent categories (code, profile,
+event) so downstream agents only receive the sub-queries relevant to them.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from langchain_core.language_models import BaseChatModel
+
+from src.agents.base import BaseAgent
+from src.prompts.classifier import build_system_prompt, pack_classification_query
+from src.schemas.classification import ClassificationResult
+from src.utils.text import parse_raw_response_to_classifications
+
+
+class ClassifierAgent(BaseAgent):
+    def __init__(self, model: BaseChatModel) -> None:
+        super().__init__(
+            model=model,
+            name="classifier",
+            system_prompt=build_system_prompt(),
+        )
+
+    async def arun(self, state: Dict[str, Any]) -> ClassificationResult:
+        user_input = state.get("user_query")
+        if not user_input:
+            self.logger.debug("Empty query — returning empty classifications.")
+            return ClassificationResult(classifications=[])
+
+        user_message = pack_classification_query(user_input)
+        messages = self._build_messages(user_message)
+        raw_content = await self._call_model(messages)
+        classifications = parse_raw_response_to_classifications(raw_content)
+
+        if classifications:
+            self.logger.info("=" * 50)
+            self.logger.info("Extracted Classifications:")
+            for idx, cls in enumerate(classifications, 1):
+                self.logger.info(
+                    "  %d. source=%s  query=%s", idx, cls["source"], cls["query"]
+                )
+            self.logger.info("Total classifications: %d", len(classifications))
+            self.logger.info("=" * 50)
+        else:
+            self.logger.info("No actionable classifications found (trivial input).")
+
+        return ClassificationResult(classifications=classifications)

src/config/constants.py

@@ -0,0 +1,12 @@
+"""
+Shared constants used across Xmem agents and prompt formatting.
+
+These are protocol-level values that all agents rely on for structured
+communication with the LLM. Changing them requires updating every
+system prompt that references the separator format.
+"""
+
+# Delimiter used in the tab-separated format between LLM and agents.
+# Format in prompts:  `- SOURCE::QUERY`
+# Must stay in sync with all system prompts and parsing utilities.
+LLM_TAB_SEPARATOR: str = "::"

src/config/settings.py

@@ -1,4 +1,4 @@
-from typing import Optional,List
+from typing import Optional, List
 from pydantic import Field,field_validator
 from pydantic_settings import BaseSettings,SettingsConfigDict
 
@@ -122,7 +122,6 @@ class Settings(BaseSettings):
     @field_validator("fallback_order")
     @classmethod
     def validate_fallback_order(cls, v: List[str]) -> List[str]:
-        """Ensure fallback_order only contains valid provider names."""
         valid_providers = {"gemini", "claude", "openai"}
         for provider in v:
             if provider not in valid_providers:

src/models/__init__.py

@@ -0,0 +1,6 @@
+"""Xmem models — re-export the public API."""
+
+from src.models.base import Provider
+from src.models.registry import get_model
+
+__all__ = ["get_model", "Provider"]

src/models/base.py

@@ -0,0 +1,7 @@
+"""
+Base types for the models module.
+"""
+
+from typing import Literal
+
+Provider = Literal["gemini", "claude", "openai"]

src/models/claude.py

@@ -0,0 +1,23 @@
+"""
+Claude model factory.
+"""
+
+from langchain_anthropic import ChatAnthropic
+from langchain_core.language_models import BaseChatModel
+
+from src.config import settings
+
+
+def build_claude_model(
+    model_name: str | None = None,
+    temperature: float | None = None,
+) -> BaseChatModel:
+    api_key = settings.claude_api_key
+    if not api_key:
+        raise ValueError("CLAUDE_API_KEY is not set")
+
+    return ChatAnthropic(
+        model=model_name or settings.claude_model,
+        api_key=api_key,
+        temperature=temperature if temperature is not None else settings.temperature,
+    )