traia-iatp 0.1.2__py3-none-any.whl → 0.1.67__py3-none-any.whl

traia_iatp/registry/iatp_search_api.py

@@ -15,15 +15,23 @@ import logging
 from pymongo import MongoClient
 from pymongo import server_api
 
-# Import for embeddings
-from .embeddings import get_embedding_service
-
 # Get environment variables
 CLUSTER_URI = "traia-iatp-cluster.yzwjvgd.mongodb.net/?retryWrites=true&w=majority&appName=Traia-IATP-Cluster"
 DATABASE_NAME = "iatp"
 
 logger = logging.getLogger(__name__)
 
+# Projection to exclude large embedding fields from query responses
+# These fields are used internally for vector search but not needed by clients
+MCP_PROJECTION_EXCLUDE_EMBEDDINGS = {
+    "description_embedding": 0,
+    "capabilities_embedding": 0
+}
+
+UTILITY_AGENT_PROJECTION_EXCLUDE_EMBEDDINGS = {
+    "embeddings": 0  # Excludes embeddings.description, embeddings.agent_card, embeddings.search_text
+}
+
 
 @dataclass
 class MCPServerInfo:
@@ -36,10 +44,25 @@ class MCPServerInfo:
     capabilities: List[str]
     metadata: Dict[str, Any]
     tags: List[str]
+    is_active: bool = False
+    core_tests_passed: bool = False
+    crewai_tests_passed: bool = False
+    user_uuid: Optional[str] = None
+    # Blockchain fields (from PostgreSQL, synced to MongoDB)
+    operator_address: Optional[str] = None
+    server_address: Optional[str] = None
+    server_network: Optional[str] = None
+    server_chain_id: Optional[int] = None
+    # Token acceptance information
+    accepts: Optional[Dict[str, Any]] = None
     
     @classmethod
     def from_registry_doc(cls, doc: Dict[str, Any]) -> 'MCPServerInfo':
-        """Create MCPServerInfo from MongoDB document."""
+        """Create MCPServerInfo from MongoDB document.
+        
+        Note: Embedding fields (description_embedding, capabilities_embedding) 
+        should be excluded from queries via projection - they are not needed by clients.
+        """
         # Extract tags from metadata if present
         metadata = doc.get('metadata', {})
         tags = metadata.get('tags', [])
@@ -52,7 +75,18 @@ class MCPServerInfo:
             server_type=doc.get('server_type', 'streamable-http'),
             capabilities=doc.get('capabilities', []),
             metadata=metadata,
-            tags=tags
+            tags=tags,
+            is_active=doc.get('is_active', False),
+            core_tests_passed=doc.get('core_tests_passed', False),
+            crewai_tests_passed=doc.get('crewai_tests_passed', False),
+            user_uuid=doc.get('user_uuid'),
+            # Blockchain fields
+            operator_address=doc.get('operator_address'),
+            server_address=doc.get('server_address'),
+            server_network=doc.get('server_network'),
+            server_chain_id=doc.get('server_chain_id'),
+            # Token acceptance
+            accepts=doc.get('accepts')
         )
 
 
@@ -69,10 +103,18 @@ class UtilityAgentInfo:
     metadata: Dict[str, Any]
     skills: List[Dict[str, Any]]
     endpoints: Optional[Dict[str, Any]] = None
+    # D402 payment fields
+    d402_enabled: bool = False
+    d402_payment_info: Optional[Dict[str, Any]] = None
     
     @classmethod
     def from_registry_doc(cls, doc: Dict[str, Any]) -> 'UtilityAgentInfo':
-        """Create UtilityAgentInfo from MongoDB document."""
+        """Create UtilityAgentInfo from MongoDB document.
+        
+        Note: The "embeddings" object (containing embeddings.description, 
+        embeddings.agent_card, embeddings.search_text) should be excluded from 
+        queries via projection - they are not needed by clients.
+        """
         return cls(
             agent_id=doc.get('agent_id', ''),
             name=doc.get('name', ''),
@@ -83,13 +125,21 @@ class UtilityAgentInfo:
             is_active=doc.get('is_active', True),
             metadata=doc.get('metadata', {}),
             skills=doc.get('skills', []),
-            endpoints=doc.get('endpoints')  # Get endpoints from root level
+            endpoints=doc.get('endpoints'),  # Get endpoints from root level
+            d402_enabled=doc.get('d402_enabled', False),
+            d402_payment_info=doc.get('d402_payment_info')
         )
 
 
 def get_readonly_connection_string() -> str:
     """Get read-only MongoDB connection string."""
-    # Try X.509 certificate authentication first
+
+    # Try IAM access
+    cluster_host_name = os.environ.get('CLUSTER_HOST_NAME', None)
+    if cluster_host_name:
+        return f"mongodb+srv://{cluster_host_name}/?authMechanism=MONGODB-AWS&authSource=%24external&retryWrites=true&w=majority&appName=Traia-IATP-Cluster"
+
+    # Try X.509 certificate authentication next
     cert_file = os.getenv("MONGODB_X509_CERT_FILE")
     if cert_file and os.path.exists(cert_file):
         # Extract just the cluster hostname without query parameters
@@ -110,9 +160,10 @@ def get_readonly_connection_string() -> str:
     
     raise ValueError(
         "MongoDB authentication required. Please provide either:\n"
-        "1. MONGODB_X509_CERT_FILE - Path to X.509 certificate file\n"
-        "2. MONGODB_USER and MONGODB_PASSWORD - Username and password\n"
-        "3. MONGODB_CONNECTION_STRING - Full connection string"
+        "1. MONGODB_IAM_ACCESS - Cluster Host Required\n"
+        "2. MONGODB_X509_CERT_FILE - Path to X.509 certificate file\n"
+        "3. MONGODB_USER and MONGODB_PASSWORD - Username and password\n"
+        "4. MONGODB_CONNECTION_STRING - Full connection string"
     )
 
 
@@ -178,14 +229,20 @@ class IATPSearchAPI:
         collection = db[collection_names["utility_agent"]]
         
         if agent_id:
-            # Direct lookup by agent_id
-            doc = collection.find_one({"agent_id": agent_id, "is_active": True})
+            # Direct lookup by agent_id (exclude embeddings for performance)
+            doc = collection.find_one(
+                {"agent_id": agent_id, "is_active": True},
+                UTILITY_AGENT_PROJECTION_EXCLUDE_EMBEDDINGS
+            )
             if doc:
                 return UtilityAgentInfo.from_registry_doc(doc)
         
         if name:
-            # Direct lookup by name
-            doc = collection.find_one({"name": name, "is_active": True})
+            # Direct lookup by name (exclude embeddings for performance)
+            doc = collection.find_one(
+                {"name": name, "is_active": True},
+                UTILITY_AGENT_PROJECTION_EXCLUDE_EMBEDDINGS
+            )
             if doc:
                 return UtilityAgentInfo.from_registry_doc(doc)
         
@@ -232,8 +289,8 @@ class IATPSearchAPI:
             if tag:
                 filters["tags"] = tag
             
-            # Find first matching document
-            doc = collection.find_one(filters)
+            # Find first matching document (exclude embeddings for performance)
+            doc = collection.find_one(filters, UTILITY_AGENT_PROJECTION_EXCLUDE_EMBEDDINGS)
             if doc:
                 return UtilityAgentInfo.from_registry_doc(doc)
         
@@ -272,8 +329,11 @@ class IATPSearchAPI:
         if capabilities:
             filters["capabilities"] = {"$in": capabilities}
         
-        # Query with limit
-        cursor = collection.find(filters).limit(limit).sort("registered_at", -1)
+        # Query with limit (exclude embeddings for performance)
+        cursor = collection.find(
+            filters,
+            UTILITY_AGENT_PROJECTION_EXCLUDE_EMBEDDINGS
+        ).limit(limit).sort("registered_at", -1)
         
         return [UtilityAgentInfo.from_registry_doc(doc) for doc in cursor]
     
@@ -298,8 +358,9 @@ class IATPSearchAPI:
         Returns:
             List of UtilityAgentInfo objects
         """
-        # Generate query embedding
+        # Generate query embedding (lazy import to avoid loading OpenAI/Cohere unless needed)
         try:
+            from .embeddings import get_embedding_service
             embedding_service = get_embedding_service()
             query_embedding = await embedding_service.generate_query_embedding(query)
         except Exception as e:
@@ -469,7 +530,7 @@ class IATPSearchAPI:
         
         if name:
             # Direct lookup by name
-            doc = collection.find_one({"name": name, "is_active": True})
+            doc = collection.find_one({"name": name, "is_active": True, "core_tests_passed": True})
             if doc:
                 return MCPServerInfo.from_registry_doc(doc)
         
@@ -488,7 +549,7 @@ class IATPSearchAPI:
                         }
                     }
                 },
-                {"$match": {"is_active": True}},
+                {"$match": {"is_active": True, "core_tests_passed": True}},
             ]
             
             # Add additional filters if provided
@@ -502,6 +563,8 @@ class IATPSearchAPI:
             if match_filters:
                 pipeline.append({"$match": match_filters})
             
+            # Exclude embedding fields from results
+            pipeline.append({"$project": MCP_PROJECTION_EXCLUDE_EMBEDDINGS})
             pipeline.append({"$limit": 1})
             
             results = list(collection.aggregate(pipeline))
@@ -509,7 +572,7 @@ class IATPSearchAPI:
                 return MCPServerInfo.from_registry_doc(results[0])
         else:
             # Build standard query filters
-            filters = {"is_active": True}
+            filters = {"is_active": True, "core_tests_passed": True}
             
             if capability:
                 filters["capabilities"] = capability
@@ -518,8 +581,8 @@ class IATPSearchAPI:
                 # Tags are stored in metadata
                 filters["metadata.tags"] = tag
             
-            # Find first matching document
-            doc = collection.find_one(filters)
+            # Find first matching document (exclude embeddings for performance)
+            doc = collection.find_one(filters, MCP_PROJECTION_EXCLUDE_EMBEDDINGS)
             if doc:
                 return MCPServerInfo.from_registry_doc(doc)
         
@@ -529,31 +592,58 @@ class IATPSearchAPI:
     def list_mcp_servers(
         cls,
         limit: int = 10,
+        page: int = 1,
         tags: Optional[List[str]] = None,
-        capabilities: Optional[List[str]] = None
+        capabilities: Optional[List[str]] = None,
+        user_uuid: Optional[str] = None
     ) -> List[MCPServerInfo]:
         """
         List available MCP servers from the registry.
         
         Args:
-            limit: Maximum number of servers to return
+            limit: Maximum number of servers to return (default: 10)
+            page: Page number for pagination, 1-indexed (default: 1)
             tags: Filter by tags
             capabilities: Filter by capabilities
+            user_uuid: Filter by user UUID
             
         Returns:
             List of MCPServerInfo objects
+            
+        Examples:
+            list_mcp_servers(limit=10, page=1)  # First 10 servers
+            list_mcp_servers(limit=10, page=2)  # Next 10 servers
         """
+        if page < 1:
+            page = 1
+        
+        skip = (page - 1) * limit
+        logger.info(f"list_mcp_servers: limit={limit}, page={page}, tags={tags}, capabilities={capabilities}, user_uuid={user_uuid}")
+        
         db = cls._get_connection()
         collection_names = get_collection_names()
         collection = db[collection_names["mcp_server"]]
         
+        logger.info(f"Using collection: {collection_names['mcp_server']}")
+        
         # Build query filters
-        filters = {"is_active": True}
+        filters = {"is_active": True, "core_tests_passed": True}
         if capabilities:
             filters["capabilities"] = {"$in": capabilities}
+        if user_uuid:
+            filters["user_uuid"] = user_uuid
         
-        # Query
-        cursor = collection.find(filters).limit(limit).sort("registered_at", -1)
+        logger.info(f"Query filters: {filters}")
+        
+        # Count and log total
+        total_count = collection.count_documents(filters)
+        logger.info(f"Total documents matching filters: {total_count}")
+        
+        # Query with pagination (exclude embeddings for performance)
+        cursor = collection.find(
+            filters, 
+            MCP_PROJECTION_EXCLUDE_EMBEDDINGS
+        ).skip(skip).limit(limit).sort("registered_at", -1)
         
         # Manual filtering for tags if needed
         servers = []
@@ -562,9 +652,11 @@ class IATPSearchAPI:
                 doc_tags = doc.get("metadata", {}).get("tags", [])
                 if not any(tag in doc_tags for tag in tags):
                     continue
+            
             servers.append(MCPServerInfo.from_registry_doc(doc))
         
-        return servers[:limit]
+        logger.info(f"Returning {len(servers)} servers")
+        return servers
     
     @classmethod
     def get_mcp_server(cls, name: str) -> Optional[Dict[str, Any]]:
@@ -581,7 +673,8 @@ class IATPSearchAPI:
         collection_names = get_collection_names()
         collection = db[collection_names["mcp_server"]]
         
-        doc = collection.find_one({"name": name})
+        # Exclude embeddings from response (not needed by clients)
+        doc = collection.find_one({"name": name}, MCP_PROJECTION_EXCLUDE_EMBEDDINGS)
         if doc:
             doc["_id"] = str(doc["_id"])
         return doc
@@ -615,8 +708,9 @@ class IATPSearchAPI:
         Returns:
             List of MCPServerInfo objects
         """
-        # Generate query embedding
+        # Generate query embedding (lazy import to avoid loading OpenAI/Cohere unless needed)
         try:
+            from .embeddings import get_embedding_service
             embedding_service = get_embedding_service()
             query_embedding = await embedding_service.generate_query_embedding(query)
         except Exception as e:
@@ -652,7 +746,7 @@ class IATPSearchAPI:
                         "queryVector": query_embedding,
                         "numCandidates": limit * 5,
                         "limit": limit,
-                        "filter": {"is_active": True} if active_only else {}
+                        "filter": {"is_active": True, "core_tests_passed": True} if active_only else {"core_tests_passed": True}
                     }
                 },
                 {
@@ -666,6 +760,14 @@ class IATPSearchAPI:
                         "metadata": 1,
                         "registered_at": 1,
                         "is_active": 1,
+                        "core_tests_passed": 1,
+                        "crewai_tests_passed": 1,
+                        "user_uuid": 1,
+                        "operator_address": 1,
+                        "server_address": 1,
+                        "server_network": 1,
+                        "server_chain_id": 1,
+                        "accepts": 1,
                         "score": {"$meta": "vectorSearchScore"}
                     }
                 }
@@ -687,7 +789,7 @@ class IATPSearchAPI:
                         "queryVector": query_embedding,
                         "numCandidates": limit * 5,
                         "limit": limit,
-                        "filter": {"is_active": True} if active_only else {}
+                        "filter": {"is_active": True, "core_tests_passed": True} if active_only else {"core_tests_passed": True}
                     }
                 },
                 {
@@ -701,6 +803,14 @@ class IATPSearchAPI:
                         "metadata": 1,
                         "registered_at": 1,
                         "is_active": 1,
+                        "core_tests_passed": 1,
+                        "crewai_tests_passed": 1,
+                        "user_uuid": 1,
+                        "operator_address": 1,
+                        "server_address": 1,
+                        "server_network": 1,
+                        "server_chain_id": 1,
+                        "accepts": 1,
                         "score": {"$meta": "vectorSearchScore"}
                     }
                 }
@@ -749,7 +859,9 @@ class IATPSearchAPI:
         ]
         
         if active_only:
-            pipeline.append({"$match": {"is_active": True}})
+            pipeline.append({"$match": {"is_active": True, "core_tests_passed": True}})
+        else:
+            pipeline.append({"$match": {"core_tests_passed": True}})
         
         pipeline.append({"$limit": limit})
         
@@ -807,11 +919,29 @@ def find_mcp_server(
 
 def list_mcp_servers(
     limit: int = 10,
+    page: int = 1,
     tags: Optional[List[str]] = None,
-    capabilities: Optional[List[str]] = None
+    capabilities: Optional[List[str]] = None,
+    user_uuid: Optional[str] = None
 ) -> List[MCPServerInfo]:
-    """List available MCP servers from the registry."""
-    return IATPSearchAPI.list_mcp_servers(limit, tags, capabilities)
+    """
+    List available MCP servers from the registry.
+    
+    Args:
+        limit: Maximum number of servers to return (default: 10)
+        page: Page number for pagination, 1-indexed (default: 1)
+        tags: Filter by tags
+        capabilities: Filter by capabilities
+        user_uuid: Filter by user UUID
+        
+    Returns:
+        List of MCPServerInfo objects
+        
+    Examples:
+        servers = list_mcp_servers(limit=10, page=1)  # First page
+        servers = list_mcp_servers(limit=10, page=2)  # Second page
+    """
+    return IATPSearchAPI.list_mcp_servers(limit, page, tags, capabilities, user_uuid)
 
 
 async def search_mcp_servers(