backend/memory/knowledge_index.py

import json
import logging
from typing import List, Dict, Any, Tuple
from memory.embedding_service import EmbeddingService
from memory.vector_store import VectorStore

logger = logging.getLogger("knowledge_index")

class KnowledgeIndexBuilder:
    def __init__(self, embedding_service: EmbeddingService, vector_store: VectorStore):
        self.embedder = embedding_service
        self.store = vector_store

    def chunk_text(self, text: str, chunk_size: int = 1000, overlap: int = 100) -> List[str]:
        """Simple sliding window text splitter."""
        if not text:
            return []
        chunks = []
        start = 0
        text_len = len(text)
        
        while start < text_len:
            end = min(start + chunk_size, text_len)
            chunks.append(text[start:end])
            start += chunk_size - overlap
            
            # Prevent infinite loop if overlap >= chunk_size
            if chunk_size - overlap <= 0:
                break
                
        return chunks

    def build_index(
        self,
        repo_id: str,
        profile: Dict[str, Any],
        summary: Dict[str, Any],
        graph: Dict[str, Any],
        report: str,
        flat_files: List[Dict[str, Any]]
    ):
        """Builds and indexes a repository's code, structure, and profile metadata into ChromaDB."""
        logger.info(f"Starting index build for repository: {repo_id}")
        
        # Clear existing collection if any
        self.store.delete_collection(repo_id)
        
        documents: List[str] = []
        metadatas: List[Dict[str, Any]] = []
        ids: List[str] = []
        
        # Helper to generate unique IDs
        def add_chunk(content: str, metadata: Dict[str, Any], prefix: str):
            chunk_id = f"{prefix}_{len(documents)}"
            documents.append(content)
            metadatas.append(metadata)
            ids.append(chunk_id)

        # 1. Index the Repository Intelligence Report (Markdown)
        report_chunks = self.chunk_text(report, chunk_size=1000, overlap=100)
        for i, chunk in enumerate(report_chunks):
            add_chunk(
                content=chunk,
                metadata={"category": "report", "chunk_index": i},
                prefix="report"
            )

        # 2. Index the Summary details
        elevator_pitch = summary.get("elevator_pitch", "")
        if elevator_pitch:
            add_chunk(
                content=f"Project Elevator Pitch:\n{elevator_pitch}",
                metadata={"category": "summary", "subcategory": "elevator_pitch"},
                prefix="summary_pitch"
            )
            
        for feat in summary.get("core_features", []):
            add_chunk(
                content=f"Core Feature: {feat}",
                metadata={"category": "summary", "subcategory": "core_feature"},
                prefix="summary_feat"
            )
            
        for start_point in summary.get("developer_start_points", []):
            add_chunk(
                content=f"Developer Starting Point File: {start_point}",
                metadata={"category": "summary", "subcategory": "developer_start_point"},
                prefix="summary_start"
            )

        # 3. Index Profile Metadata (Architecture, APIs, Auth, Dependencies)
        arch_pattern = profile.get("architecture_pattern", "")
        if arch_pattern:
            add_chunk(
                content=f"Architecture Pattern: {arch_pattern}",
                metadata={"category": "architecture", "pattern": arch_pattern},
                prefix="profile_arch"
            )
            
        for auth_method in profile.get("authentication_methods", []):
            add_chunk(
                content=f"Authentication / Security Method: {auth_method}",
                metadata={"category": "authentication", "method": auth_method},
                prefix="profile_auth"
            )
            
        for endpoint in profile.get("api_endpoints", []):
            add_chunk(
                content=f"API Endpoint / Route: {endpoint}",
                metadata={"category": "api", "endpoint": endpoint},
                prefix="profile_api"
            )
            
        for dep in profile.get("dependencies", []):
            add_chunk(
                content=f"Package Dependency: {dep}",
                metadata={"category": "dependency", "dependency": dep},
                prefix="profile_dep"
            )

        # 4. Index Business Flows & Concepts from Graph
        for flow in graph.get("business_flows", []):
            flow_name = flow.get("flow_name", "")
            flow_desc = flow.get("description", "")
            flow_steps = ", ".join(flow.get("steps", []))
            add_chunk(
                content=f"Business Flow: {flow_name}\nDescription: {flow_desc}\nSteps: {flow_steps}",
                metadata={"category": "business_flow", "flow_name": flow_name},
                prefix="graph_flow"
            )
            
        for concept in graph.get("concepts", []):
            concept_name = concept.get("name", "")
            concept_desc = concept.get("description", "")
            concept_files = ", ".join(concept.get("files", []))
            add_chunk(
                content=f"Core Concept: {concept_name}\nDescription: {concept_desc}\nFiles: {concept_files}",
                metadata={"category": "concept", "concept_name": concept_name},
                prefix="graph_concept"
            )

        # 5. Index Source Code Files Content
        for f in flat_files:
            file_path = f.get("path", "")
            content = f.get("content", "")
            file_size = f.get("size", 0)
            
            # Skip massive files to avoid cluttering vector space
            if file_size > 100 * 1024 or not content:
                continue
                
            file_chunks = self.chunk_text(content, chunk_size=1000, overlap=100)
            for j, chunk in enumerate(file_chunks):
                # Include path info inside the document text to maintain retrieval association
                chunk_content = f"File: {file_path} (Chunk {j+1}/{len(file_chunks)})\n\n{chunk}"
                add_chunk(
                    content=chunk_content,
                    metadata={"category": "file", "path": file_path, "chunk_index": j},
                    prefix="file_chunk"
                )

        if not documents:
            logger.info("No documents found to index.")
            return

        # 6. Generate Embeddings in Batches of 50 to respect API rate limits and connection pooling
        batch_size = 50
        logger.info(f"Generating embeddings for {len(documents)} document chunks in batches of {batch_size}...")
        
        all_embeddings: List[List[float]] = []
        for start_idx in range(0, len(documents), batch_size):
            end_idx = min(start_idx + batch_size, len(documents))
            batch_docs = documents[start_idx:end_idx]
            batch_embeddings = self.embedder.embed_texts(batch_docs)
            all_embeddings.extend(batch_embeddings)
            
        # 7. Write to Vector Store
        self.store.add_documents(
            repo_id=repo_id,
            documents=documents,
            metadatas=metadatas,
            ids=ids,
            embeddings=all_embeddings
        )
        logger.info(f"Vector database indexing complete. Indexed {len(documents)} chunks.")