adaptive-memory-multi-model-router 2.16.0 → 2.16.2

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"),
+    )

package/adapters/a3m_adapter/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for A3M adapters."""

package/adapters/a3m_adapter/tests/test_adapters.py

@@ -0,0 +1,118 @@
+"""
+Test script for A3M Router adapters.
+
+Tests the LangChain and LlamaIndex adapters to ensure they:
+1. Initialize correctly
+2. Route requests properly
+3. Return expected response types
+4. Handle errors gracefully
+"""
+
+import sys
+import os
+import logging
+
+# Add current directory to path
+sys.path.insert(0, '.')
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def test_langchain_adapter():
+    """Test LangChain adapter."""
+    print("Testing LangChain adapter...")
+    
+    try:
+        from a3m_adapter import A3MLangChainAdapter
+        
+        # Initialize
+        llm = A3MLangChainAdapter(model="auto", temperature=0.7)
+        print(f"✅ Initialized: {llm}")
+        
+        # Test simple generation
+        # Note: This would make actual API calls - we'll skip for now
+        # In a real test, we'd mock the A3M router
+        print("✅ LangChain adapter structure OK")
+        assert llm is not None
+        
+    except Exception as e:
+        print(f"❌ LangChain adapter failed: {e}")
+        raise
+
+
+def test_llamaindex_adapter():
+    """Test LlamaIndex adapter."""
+    print("Testing LlamaIndex adapter...")
+    
+    try:
+        from a3m_adapter import A3MLlamaIndexAdapter
+        
+        # Initialize
+        llm = A3MLlamaIndexAdapter(model="auto", temperature=0.5)
+        print(f"✅ Initialized: {llm}")
+        
+        # Check metadata
+        metadata = llm.metadata
+        print(f"✅ Metadata: {metadata}")
+        assert metadata is not None
+        
+    except Exception as e:
+        print(f"❌ LlamaIndex adapter failed: {e}")
+        raise
+
+
+def test_config():
+    """Test configuration."""
+    print("Testing configuration...")
+    
+    try:
+        from a3m_adapter import A3MConfig
+        
+        # Test defaults
+        config = A3MConfig()
+        print(f"✅ Default config: model={config.model}")
+        
+        # Test to_dict
+        data = config.to_dict()
+        assert 'model' in data
+        print("✅ Config to_dict works")
+        
+        # Test JSON serialization
+        json_str = config.to_json()
+        assert '"model"' in json_str
+        print("✅ Config JSON serialization works")
+        
+    except Exception as e:
+        print(f"❌ Config test failed: {e}")
+        raise
+
+
+def main():
+    """Run all tests."""
+    print("=" * 50)
+    print("A3M Router Adapter Tests")
+    print("=" * 50)
+    
+    tests = [
+        test_config,
+        test_langchain_adapter,
+        test_llamaindex_adapter,
+    ]
+    
+    passed = 0
+    total = len(tests)
+    
+    for test in tests:
+        test()
+        passed += 1
+        print()
+    
+    print("=" * 50)
+    print(f"Results: {passed}/{total} tests passed")
+    print("🎉 All tests passed!")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/adapters/a3m_adapter/tests/test_integration.py

@@ -0,0 +1,80 @@
+"""
+Integration tests for A3M Router adapters.
+Requires A3M Router server running on localhost:8787
+"""
+
+import pytest
+import os
+
+
+@pytest.fixture
+def a3m_server_url():
+    """Get A3M Router server URL."""
+    return os.environ.get("A3M_SERVER_URL", "http://localhost:8787")
+
+
+@pytest.fixture
+def skip_if_no_server():
+    """Skip test if server is not available."""
+    import requests
+    try:
+        resp = requests.get("http://localhost:8787/health", timeout=2)
+        if resp.status_code != 200:
+            pytest.skip("A3M Router server not running")
+    except:
+        pytest.skip("A3M Router server not running")
+
+
+@pytest.mark.integration
+def test_simple_chat_completion(a3m_server_url, skip_if_no_server):
+    """Test simple chat completion via HTTP API."""
+    import requests
+    
+    response = requests.post(
+        f"{a3m_server_url}/v1/chat/completions",
+        json={
+            "model": "auto",
+            "messages": [{"role": "user", "content": "What is 2+2?"}]
+        },
+        timeout=30
+    )
+    
+    assert response.status_code == 200
+    data = response.json()
+    assert "choices" in data
+    assert len(data["choices"]) > 0
+    assert "message" in data["choices"][0]
+    assert data["choices"][0]["message"]["content"]
+
+
+@pytest.mark.integration
+def test_parallel_ensemble(a3m_server_url, skip_if_no_server):
+    """Test parallel ensemble with multiple providers."""
+    import requests
+    
+    response = requests.post(
+        f"{a3m_server_url}/v1/chat/completions",
+        json={
+            "model": "auto",
+            "messages": [{"role": "user", "content": "Explain gravity"}],
+            "parallel_ensemble": 3,
+        },
+        timeout=60
+    )
+    
+    assert response.status_code == 200
+    data = response.json()
+    assert "choices" in data
+    assert "provider" in data
+
+
+@pytest.mark.integration  
+def test_health_endpoint(a3m_server_url, skip_if_no_server):
+    """Test health endpoint."""
+    import requests
+    
+    response = requests.get(f"{a3m_server_url}/health", timeout=10)
+    
+    assert response.status_code == 200
+    data = response.json()
+    assert "providers" in data or "status" in data

package/adapters/requirements-dev.txt

@@ -0,0 +1,6 @@
+-r requirements.txt
+black>=21.12
+isort>=5.10.1
+flake8>=3.9.0
+build>=0.10.0
+twine>=4.0.0

package/adapters/requirements.txt

@@ -0,0 +1,4 @@
+requests>=2.25.1
+pydantic>=1.9.0
+pytest>=6.0
+pytest-asyncio>=0.21.0

package/adapters/setup.py

@@ -0,0 +1,23 @@
+from setuptools import setup, find_packages
+import os
+
+setup(
+    name="a3m_adapter",
+    version="1.0.0",
+    description="A3M Router adapters for LangChain, LlamaIndex, and other LLM frameworks",
+    long_description=open("README.md").read() if os.path.exists("README.md") else "",
+    long_description_content_type="text/markdown",
+    author="A3M Team",
+    author_email="hello@a3m.ai",
+    packages=find_packages(),
+    install_requires=[
+        "requests>=2.25.1",
+        "pydantic>=1.9.0",
+    ],
+    extras_require={
+        "langchain": ["langchain>=0.0.365", "langchain-core>=0.0.365"],
+        "llamaindex": ["llama-index>=0.8.0"],
+        "dev": ["pytest>=6.0"],
+    },
+    python_requires=">=3.8",
+)