""" Agent OS Ingestion Module Handles ingestion of Agent OS profiles into Claude OS knowledge bases. Integrates with ChromaManager to store parsed Agent OS content. """ import logging from typing import List, Dict, Any, Optional from pathlib import Path from llama_index.embeddings.ollama import OllamaEmbedding from app.core.agent_os_parser import AgentOSParser, AgentOSDocument, AgentOSContentType from app.core.sqlite_manager import SQLiteManager from app.core.config import Config logger = logging.getLogger(__name__) class AgentOSIngestion: """Handles ingestion of Agent OS profiles into knowledge bases.""" def __init__(self, db_manager: SQLiteManager): """ Initialize Agent OS ingestion. Args: db_manager: SQLiteManager instance for storage """ self.db_manager = db_manager self.parser = AgentOSParser() def ingest_profile( self, kb_name: str, profile_path: str, batch_size: int = 100 ) -> Dict[str, Any]: """ Ingest an Agent OS profile into a knowledge base. Args: kb_name: Name of the knowledge base profile_path: Path to the Agent OS profile directory batch_size: Number of documents to process in each batch Returns: Dictionary with ingestion statistics Raises: ValueError: If KB doesn't exist or profile path is invalid """ logger.info(f"Starting Agent OS profile ingestion: {profile_path} -> {kb_name}") # Validate KB exists collections = self.db_manager.list_collections() kb_names = [col["name"] for col in collections] if kb_name not in kb_names: raise ValueError(f"Knowledge base '{kb_name}' does not exist") # Parse the profile directory try: documents = self.parser.parse_directory(profile_path) except Exception as e: logger.error(f"Failed to parse Agent OS profile: {e}") raise ValueError(f"Failed to parse Agent OS profile: {e}") if not documents: logger.warning("No documents found in Agent OS profile") return { "success": False, "message": "No documents found in profile", "documents_processed": 0 } # Ingest documents in batches stats = { "total_documents": len(documents), "documents_by_type": {}, "success": True, "errors": [] } # Count documents by type for doc in documents: content_type = doc.content_type.value stats["documents_by_type"][content_type] = \ stats["documents_by_type"].get(content_type, 0) + 1 # Process in batches for i in range(0, len(documents), batch_size): batch = documents[i:i + batch_size] try: self._ingest_batch(kb_name, batch) logger.info(f"Ingested batch {i//batch_size + 1}: {len(batch)} documents") except Exception as e: error_msg = f"Failed to ingest batch {i//batch_size + 1}: {e}" logger.error(error_msg) stats["errors"].append(error_msg) stats["success"] = False logger.info(f"Agent OS ingestion complete: {stats}") return stats def _ingest_batch(self, kb_name: str, documents: List[AgentOSDocument]) -> None: """ Ingest a batch of documents into the knowledge base. Args: kb_name: Name of the knowledge base documents: List of AgentOSDocument objects to ingest """ if not documents: return # Prepare data for ChromaDB texts = [] metadatas = [] ids = [] for idx, doc in enumerate(documents): # Create unique ID doc_id = f"{doc.content_type.value}_{Path(doc.file_path).stem}_{idx}" # Prepare metadata (ChromaDB only supports string, int, float, bool) metadata = { "content_type": doc.content_type.value, "title": doc.title, "file_path": doc.file_path, "source": "agent_os", } # Add document metadata (convert complex types to strings) for key, value in doc.metadata.items(): if isinstance(value, (str, int, float, bool)): metadata[key] = value elif isinstance(value, list): # Serialize lists as comma-separated strings metadata[key] = ",".join(str(v) for v in value) else: # Convert other types to string metadata[key] = str(value) texts.append(doc.content) metadatas.append(metadata) ids.append(doc_id) # Generate embeddings embed_model = OllamaEmbedding( model_name=Config.OLLAMA_EMBED_MODEL, base_url=Config.OLLAMA_HOST ) embeddings = [embed_model.get_text_embedding(text) for text in texts] # Add to SQLite self.db_manager.add_documents( kb_name=kb_name, documents=texts, embeddings=embeddings, metadatas=metadatas, ids=ids ) logger.debug(f"Added {len(documents)} documents to {kb_name}") def get_profile_stats(self, kb_name: str) -> Dict[str, Any]: """ Get statistics about Agent OS content in a knowledge base. Args: kb_name: Name of the knowledge base Returns: Dictionary with statistics """ try: # Get all documents with agent_os source results = self.db_manager.get_documents_by_metadata( kb_name=kb_name, where={"source": "agent_os"} ) if not results: return { "total_documents": 0, "documents_by_type": {} } # Count by content type type_counts = {} for doc in results: metadata = doc.get("metadata", {}) content_type = metadata.get("content_type", "unknown") type_counts[content_type] = type_counts.get(content_type, 0) + 1 return { "total_documents": len(results), "documents_by_type": type_counts } except Exception as e: logger.error(f"Failed to get profile stats for {kb_name}: {e}") return { "total_documents": 0, "documents_by_type": {}, "error": str(e) } def search_by_type( self, kb_name: str, content_type: AgentOSContentType, query: Optional[str] = None, limit: int = 10 ) -> List[Dict[str, Any]]: """ Search Agent OS content by type. Args: kb_name: Name of the knowledge base content_type: Type of content to search query: Optional search query (if None, returns all of type) limit: Maximum number of results Returns: List of matching documents """ try: where_filter = { "source": "agent_os", "content_type": content_type.value } if query: # Semantic search with type filter embed_model = OllamaEmbedding( model_name=Config.OLLAMA_EMBED_MODEL, base_url=Config.OLLAMA_HOST ) query_embedding = embed_model.get_text_embedding(query) results = self.db_manager.query_documents( kb_name=kb_name, query_embedding=query_embedding, n_results=limit, where=where_filter ) # Format query results documents = [] for i in range(len(results.get("documents", []))): documents.append({ "content": results["documents"][i], "metadata": results["metadatas"][i], "distance": results.get("distances", [0])[i] }) else: # Just get all of this type results = self.db_manager.get_documents_by_metadata( kb_name=kb_name, where=where_filter, limit=limit ) # Format get results documents = [] for doc in results: documents.append({ "content": doc.get("content", ""), "metadata": doc.get("metadata", {}) }) return documents except Exception as e: logger.error(f"Failed to search by type in {kb_name}: {e}") return []