adaptive-memory-multi-model-router 2.15.4 → 2.15.5

package/adapters/a3m_adapter/adapter/langgraph.py

@@ -0,0 +1,196 @@
+"""
+A3M Router Adapter for LangGraph (Microsoft's agent framework).
+
+Drop-in replacement for LangGraph's stateful agent that routes through A3M Router
+for intelligent, cost-optimized multi-step conversations.
+
+Usage:
+    from langgraph.prebuilt import create_react_agent
+    from a3m_adapter import A3MLangGraphAdapter
+    
+    adapter = A3MLangGraphAdapter(model='auto', parallel_ensemble=2)
+    
+    agent = create_react_agent(adapter, tools=[...])
+    
+    result = agent.invoke({"messages": [{"role": "user", "content": "Hello"]})
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, TypedDict
+
+logger = logging.getLogger(__name__)
+
+LANGGRAPH_AVAILABLE = False
+try:
+    import langgraph
+    from langgraph.prebuilt import create_react_agent
+    from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
+    LANGGRAPH_AVAILABLE = True
+except ImportError:
+    logger.warning("LangGraph not installed. Install with: pip install langgraph")
+
+A3M_AVAILABLE = False
+try:
+    from a3m.router import A3MRouter, RouteResponse
+    A3M_AVAILABLE = True
+except ImportError:
+    logger.warning(
+        "A3M Router not installed. Install with: pip install adaptive-memory-multi-model-router"
+    )
+
+
+class A3MLangGraphAdapter:
+    """
+    A3M Router adapter for LangGraph's prebuilt agents.
+
+    Enables LangGraph agents to use A3M Router for automatic model selection
+    across 47+ providers with cost optimization and stateful conversations.
+    """
+
+    def __init__(
+        self,
+        model: str = "auto",
+        temperature: float = 0.7,
+        max_tokens: int = 4096,
+        parallel_ensemble: int = 1,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize A3M Router adapter for LangGraph.
+        """
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.parallel_ensemble = parallel_ensemble
+        self.api_key = api_key
+        self._a3m_router = None
+        self._initialized = False
+        self._kwargs = kwargs
+
+    def _ensure_router(self) -> None:
+        """Lazily initialize the A3M router."""
+        if self._initialized:
+            return
+        
+        if not A3M_AVAILABLE:
+            raise ImportError(
+                "A3M Router is not installed. "
+                "Install with: pip install adaptive-memory-multi-model-router"
+            )
+        
+        self._a3m_router = A3MRouter(
+            model=self.model,
+            temperature=self.temperature,
+            parallel_ensemble=self.parallel_ensemble,
+        )
+        self._initialized = True
+        logger.info(
+            "A3M Router initialized for LangGraph: model=%s, ensemble=%d",
+            self.model,
+            self.parallel_ensemble,
+        )
+
+    def get_model(self):
+        """
+        Get the underlying model for LangGraph.
+        
+        Returns an object compatible with LangGraph's prebuilt agents.
+        """
+        self._ensure_router()
+        return self
+
+    def __call__(
+        self,
+        state: Dict[str, Any],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        LangGraph-compatible callable for node execution.
+        
+        Args:
+            state: LangGraph state dict with 'messages' key
+            
+        Returns:
+            Updated state dict
+        """
+        self._ensure_router()
+        
+        messages = state.get("messages", [])
+        
+        # Convert LangGraph messages to A3M format
+        a3m_messages = self._convert_messages(messages)
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=a3m_messages,
+                temperature=kwargs.get("temperature", self.temperature),
+                max_tokens=kwargs.get("max_tokens", self.max_tokens),
+                **kwargs,
+            ),
+        )
+        
+        # Add response to messages
+        new_messages = messages + [
+            AIMessage(content=route_result.content)
+        ]
+        
+        return {
+            **state,
+            "messages": new_messages,
+        }
+
+    def _convert_messages(
+        self,
+        messages: List[BaseMessage],
+    ) -> List[Dict[str, Any]]:
+        """Convert LangGraph messages to A3M format."""
+        a3m_messages = []
+        for msg in messages:
+            if isinstance(msg, HumanMessage):
+                a3m_messages.append({"role": "user", "content": msg.content})
+            elif isinstance(msg, AIMessage):
+                a3m_messages.append({"role": "assistant", "content": msg.content})
+            else:
+                a3m_messages.append({"role": "user", "content": str(msg)})
+        return a3m_messages
+
+    async def ainvoke(
+        self,
+        state: Dict[str, Any],
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Async version of __call__."""
+        self._ensure_router()
+        
+        messages = state.get("messages", [])
+        a3m_messages = self._convert_messages(messages)
+        
+        route_result = await self._a3m_router.aroute(
+            messages=a3m_messages,
+            temperature=kwargs.get("temperature", self.temperature),
+            max_tokens=kwargs.get("max_tokens", self.max_tokens),
+            **kwargs,
+        )
+        
+        new_messages = messages + [
+            AIMessage(content=route_result.content)
+        ]
+        
+        return {
+            **state,
+            "messages": new_messages,
+        }
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MLangGraphAdapter("
+            f"model={self.model!r}, "
+            f"temperature={self.temperature}, "
+            f"ensemble={self.parallel_ensemble})"
+        )

package/adapters/a3m_adapter/adapter/pinecone.py

@@ -0,0 +1,217 @@
+"""
+A3M Router Adapter for Pinecone Vector Database.
+
+Enables Pinecone's managed vector database to use A3M Router for
+intelligent query routing and cost-optimized embeddings.
+
+Usage:
+    from pinecone import Pinecone
+    from a3m_adapter import A3MPineconeAdapter
+    
+    # Create A3M-powered embeddings
+    embed_adapter = A3MPineconeAdapter(model="auto")
+    
+    # Generate embeddings
+    embedding = embed_adapter.embed("What is quantum computing?")
+    
+    # Query Pinecone
+    results = index.query(
+        vector=embedding,
+        top_k=5,
+    )
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+A3M_AVAILABLE = False
+try:
+    from a3m.router import A3MRouter
+    A3M_AVAILABLE = True
+except ImportError:
+    logger.warning(
+        "A3M Router not installed. Install with: pip install adaptive-memory-multi-model-router"
+    )
+
+
+class A3MPineconeAdapter:
+    """
+    A3M Router adapter for Pinecone embeddings.
+
+    Provides intelligent embedding generation through A3M Router
+    with automatic model selection for cost optimization.
+    """
+
+    def __init__(
+        self,
+        model: str = "auto",
+        embed_model: str = "auto",
+        parallel_ensemble: int = 1,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize A3M Router adapter for Pinecone.
+        """
+        self.model = model
+        self.embed_model = embed_model or "auto"
+        self.parallel_ensemble = parallel_ensemble
+        self.api_key = api_key
+        self._a3m_router = None
+        self._initialized = False
+        self._kwargs = kwargs
+
+    def _ensure_router(self) -> None:
+        """Lazily initialize the A3M router."""
+        if self._initialized:
+            return
+        
+        if not A3M_AVAILABLE:
+            raise ImportError(
+                "A3M Router is not installed. "
+                "Install with: pip install adaptive-memory-multi-model-router"
+            )
+        
+        self._a3m_router = A3MRouter(
+            model=self.model,
+            parallel_ensemble=self.parallel_ensemble,
+        )
+        self._initialized = True
+        logger.info(
+            "A3M Router initialized for Pinecone: embed_model=%s",
+            self.embed_model,
+        )
+
+    def embed(
+        self,
+        texts: List[str],
+        **kwargs: Any,
+    ) -> List[List[float]]:
+        """
+        Generate embeddings for texts using A3M Router.
+        
+        Args:
+            texts: List of text strings to embed
+            
+        Returns:
+            List of embedding vectors
+        """
+        self._ensure_router()
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        
+        # For embeddings, we typically call the router with a special embedding mode
+        # Since A3M supports /v1/embeddings endpoint
+        results = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.embed(
+                texts=texts,
+                model=self.embed_model,
+                **kwargs,
+            ),
+        )
+        
+        return results
+
+    def embed_query(
+        self,
+        text: str,
+        **kwargs: Any,
+    ) -> List[float]:
+        """
+        Generate embedding for a single query.
+        
+        Args:
+            text: Text to embed
+            
+        Returns:
+            Embedding vector
+        """
+        embeddings = self.embed([text], **kwargs)
+        return embeddings[0] if embeddings else []
+
+    async def aembed(
+        self,
+        texts: List[str],
+        **kwargs: Any,
+    ) -> List[List[float]]:
+        """Async version of embed."""
+        self._ensure_router()
+        
+        results = await self._a3m_router.aembed(
+            texts=texts,
+            model=self.embed_model,
+            **kwargs,
+        )
+        
+        return results
+
+    def rag_query(
+        self,
+        query: str,
+        index,
+        top_k: int = 5,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Perform RAG query: embed + Pinecone search + context.
+        
+        Args:
+            query: The search query
+            index: Pinecone index to query
+            top_k: Number of results to retrieve
+            
+        Returns:
+            Dict with 'results', 'context', 'provider', 'cost'
+        """
+        # 1. Embed query
+        query_embedding = self.embed_query(query)
+        
+        # 2. Search Pinecone
+        search_results = index.query(
+            vector=query_embedding,
+            top_k=top_k,
+            include_metadata=True,
+        )
+        
+        # 3. Build context from results
+        context = "\n\n".join([
+            match.get('metadata', {}).get('text', str(match.get('id', '')))
+            for match in search_results.get('matches', [])[:3]
+        ])
+        
+        # 4. Route the full query through A3M
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=[{
+                    "role": "user",
+                    "content": f"Context:\n{context}\n\nQuestion: {query}"
+                }],
+                temperature=self._kwargs.get("temperature", 0.7),
+                **kwargs,
+            ),
+        )
+        
+        return {
+            "results": search_results.get('matches', []),
+            "context": context,
+            "answer": route_result.content,
+            "provider": getattr(route_result, 'provider', 'a3m'),
+            "cost": getattr(route_result, 'cost', 0.0),
+        }
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MPineconeAdapter("
+            f"model={self.model!r}, "
+            f"embed_model={self.embed_model!r}, "
+            f"ensemble={self.parallel_ensemble})"
+        )

package/adapters/a3m_adapter/adapter/vercel.py

@@ -0,0 +1,188 @@
+"""
+A3M Router Adapter for Vercel AI SDK.
+
+Drop-in replacement for Vercel AI SDK's AI function that routes through A3M Router
+for intelligent, cost-optimized responses in Next.js and other JavaScript environments.
+
+Usage (JavaScript):
+    import { generateText } from 'ai';
+    import { createA3MProvider } from 'a3m-adapter/vercel';
+
+    const result = await generateText({
+        model: createA3MProvider({ model: 'auto', parallel_ensemble: 2 }),
+        prompt: 'What is the meaning of life?',
+    });
+
+Usage (Python):
+    from a3m_adapter import A3MVercelAdapter
+    
+    adapter = A3MVercelAdapter(model='auto', temperature=0.7)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+A3M_AVAILABLE = False
+try:
+    from a3m.router import A3MRouter, RouteResponse
+    A3M_AVAILABLE = True
+except ImportError:
+    logger.warning(
+        "A3M Router not installed. Install with: pip install adaptive-memory-multi-model-router"
+    )
+
+
+class A3MVercelAdapter:
+    """
+    A3M Router adapter for Vercel AI SDK compatibility.
+
+    Provides a drop-in replacement that routes through A3M Router
+    instead of calling OpenAI/Anthropic directly.
+    """
+
+    def __init__(
+        self,
+        model: str = "auto",
+        temperature: float = 0.7,
+        max_tokens: int = 4096,
+        parallel_ensemble: int = 1,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize A3M Router adapter for Vercel AI SDK.
+        """
+        self.model = model
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+        self.parallel_ensemble = parallel_ensemble
+        self.api_key = api_key
+        self._a3m_router = None
+        self._initialized = False
+        self._kwargs = kwargs
+
+    def _ensure_router(self) -> None:
+        """Lazily initialize the A3M router."""
+        if self._initialized:
+            return
+        
+        if not A3M_AVAILABLE:
+            raise ImportError(
+                "A3M Router is not installed. "
+                "Install with: pip install adaptive-memory-multi-model-router"
+            )
+        
+        self._a3m_router = A3MRouter(
+            model=self.model,
+            temperature=self.temperature,
+            parallel_ensemble=self.parallel_ensemble,
+        )
+        self._initialized = True
+        logger.info(
+            "A3M Router initialized for Vercel AI SDK: model=%s",
+            self.model,
+        )
+
+    def __call__(
+        self,
+        prompt: str,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """
+        Generate text from prompt (Vercel AI SDK compatible interface).
+        
+        Args:
+            prompt: The prompt string
+            
+        Returns:
+            Dict with 'text', 'provider', 'usage', 'finishReason'
+        """
+        self._ensure_router()
+        
+        messages = [{"role": "user", "content": prompt}]
+        
+        import asyncio
+        loop = asyncio.get_event_loop()
+        route_result = loop.run_in_executor(
+            None,
+            lambda: self._a3m_router.route(
+                messages=messages,
+                temperature=kwargs.get("temperature", self.temperature),
+                max_tokens=kwargs.get("max_tokens", self.max_tokens),
+                **kwargs,
+            ),
+        )
+        
+        return {
+            "text": route_result.content,
+            "provider": getattr(route_result, 'provider', 'a3m'),
+            "finishReason": getattr(route_result, 'finish_reason', 'stop'),
+            "usage": {
+                "promptTokens": getattr(route_result, 'prompt_tokens', 0),
+                "completionTokens": getattr(route_result, 'completion_tokens', 0),
+                "totalTokens": getattr(route_result, 'total_tokens', 0),
+            },
+        }
+
+    async def generate(
+        self,
+        prompt: str,
+        **kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Async generate for Vercel AI SDK."""
+        self._ensure_router()
+        
+        messages = [{"role": "user", "content": prompt}]
+        
+        route_result = await self._a3m_router.aroute(
+            messages=messages,
+            temperature=kwargs.get("temperature", self.temperature),
+            max_tokens=kwargs.get("max_tokens", self.max_tokens),
+            **kwargs,
+        )
+        
+        return {
+            "text": route_result.content,
+            "provider": getattr(route_result, 'provider', 'a3m'),
+            "finishReason": getattr(route_result, 'finish_reason', 'stop'),
+            "usage": {
+                "promptTokens": getattr(route_result, 'prompt_tokens', 0),
+                "completionTokens": getattr(route_result, 'completion_tokens', 0),
+                "totalTokens": getattr(route_result, 'total_tokens', 0),
+            },
+        }
+
+    def __repr__(self) -> str:
+        return (
+            f"A3MVercelAdapter("
+            f"model={self.model!r}, "
+            f"temperature={self.temperature}, "
+            f"max_tokens={self.max_tokens})"
+        )
+
+
+# JavaScript-compatible factory function
+def createA3MProvider(config: Dict[str, Any]) -> A3MVercelAdapter:
+    """
+    Create an A3M Provider for Vercel AI SDK (JavaScript usage).
+    
+    Usage:
+        import { generateText } from 'ai';
+        import { createA3MProvider } from 'a3m-adapter/vercel';
+        
+        const result = await generateText({
+            model: createA3MProvider({ model: 'auto', parallel_ensemble: 2 }),
+            prompt: 'What is 2+2?',
+        });
+    """
+    return A3MVercelAdapter(
+        model=config.get("model", "auto"),
+        temperature=config.get("temperature", 0.7),
+        max_tokens=config.get("max_tokens", 4096),
+        parallel_ensemble=config.get("parallel_ensemble", 1),
+        api_key=config.get("api_key"),
+    )