ragbox-core 1.0.0__py3-none-any.whl

ragbox/__init__.py

@@ -0,0 +1,175 @@
+"""
+RAGBox: Zero-Configuration Self-Building Agentic RAG System.
+
+Provides the public RAGBox class for one-line integrations.
+"""
+import asyncio
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from loguru import logger
+
+from ragbox.config.defaults import Settings
+from ragbox.core.agentic_orchestrator import AgenticOrchestrator
+from ragbox.core.chunking_engine import ChunkingEngine
+from ragbox.core.document_processor import DocumentProcessorRouter
+from ragbox.core.knowledge_graph import OptimizedKnowledgeGraph
+from ragbox.core.retrieval_fusion import RetrievalFusionEngine
+from ragbox.core.self_healing import SelfHealer
+from ragbox.models.queries import Answer
+from ragbox.utils.embeddings import EmbeddingAutoDetector
+from ragbox.utils.llm_clients import LLMAutoDetector
+from ragbox.utils.vector_stores import VectorStoreAutoDetector
+
+
+class RAGBox:
+    """
+    RAG-in-a-Box: Zero-Configuration Self-Building Agentic RAG System.
+    """
+
+    def __init__(self, document_dir: str | Path, config: Optional[Settings] = None) -> None:
+        """
+        Initialize a complete RAG system pointing to a directory.
+        
+        Args:
+            document_dir: Path to the directory containing documents.
+            config: Optional configuration overrides.
+        """
+        self.document_dir = Path(document_dir)
+        self.config = config or Settings()
+        
+        # Initialize Core Layers Auto-magically
+        self.llm_client = LLMAutoDetector.detect(self.config)
+        self.embedding_provider = EmbeddingAutoDetector.detect(self.config)
+        self.vector_store = VectorStoreAutoDetector.detect(self.config, self.embedding_provider)
+        
+        self.document_processor = DocumentProcessorRouter()
+        self.chunking_engine = ChunkingEngine(self.llm_client, self.embedding_provider)
+        self.knowledge_graph = OptimizedKnowledgeGraph(llm_client=self.llm_client)
+        
+        self.retriever = RetrievalFusionEngine(
+            vector_store=self.vector_store,
+            knowledge_graph=self.knowledge_graph,
+            embedding_provider=self.embedding_provider,
+            llm_client=self.llm_client
+        )
+        self.orchestrator = AgenticOrchestrator(
+            retriever=self.retriever,
+            llm_client=self.llm_client,
+            document_processor=self.document_processor,
+            knowledge_graph=self.knowledge_graph
+        )
+        self.self_healer = SelfHealer(
+            document_dir=self.document_dir,
+            document_processor=self.document_processor,
+            chunking_engine=self.chunking_engine,
+            vector_store=self.vector_store,
+            knowledge_graph=self.knowledge_graph
+        )
+        
+        # Start initial build and self-healing daemon
+        logger.info(f"Initializing RAGBox for {self.document_dir}")
+        self._ensure_built()
+        
+    def _ensure_built(self) -> None:
+        """Synchronously trigger the build process to guarantee readiness."""
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            
+        if loop.is_running():
+            asyncio.create_task(self.self_healer.initial_build())
+        else:
+            loop.run_until_complete(self.self_healer.initial_build())
+        
+        # Start watchdog
+        self.self_healer.start_watchdog()
+
+    def query(self, question: str) -> str:
+        """
+        Query the RAGBox system synchronously.
+        
+        Args:
+            question: The user's query.
+            
+        Returns:
+            The direct answer string.
+        """
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                raise RuntimeError("Cannot call synchronous `query` inside an active event loop. Use `aquery`.")
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            
+        answer: Answer = loop.run_until_complete(self.orchestrator.execute(question))
+        return answer.content
+        
+    async def aquery(self, question: str) -> str:
+        """
+        Async evaluation of a query.
+        
+        Args:
+            question: The user's query.
+            
+        Returns:
+            The direct answer string.
+        """
+        answer: Answer = await self.orchestrator.execute(question)
+        return answer.content
+
+    def estimate_cost(self, query: Optional[str] = None) -> str:
+        """
+        Estimate the cost of a query or indexing the corpus upfront.
+
+        Args:
+            query: The query to run. If None, estimates the cost of indexing the entire corpus.
+
+        Returns:
+            A string detailing the estimated cost.
+        """
+        from ragbox.utils.cost_tracker import CostEstimator
+        
+        # We assume the default model for estimating cost
+        estimator = CostEstimator(self.llm_client.model_name if hasattr(self.llm_client, "model_name") else "gpt-4o")
+        
+        if query is None:
+            # Estimate cost of indexing the corpus based on file size
+            total_bytes = 0
+            for ext in self.document_processor.processors.keys():
+                for p in self.document_dir.rglob(f"*{ext}"):
+                    if p.is_file():
+                        total_bytes += p.stat().st_size
+            
+            # Wild guess: 1 byte ~ 0.25 tokens
+            estimated_tokens = int(total_bytes * 0.25)
+            # Embedding cost
+            embed_cost = (estimated_tokens / 1_000_000) * 0.13 # large embedding assumption
+            return f"Corpus Indexing Estimate: ~{estimated_tokens} tokens. Embedding Cost: ${embed_cost:.4f}"
+            
+        else:
+            # Simulate retrieval and estimation for query
+            try:
+                loop = asyncio.get_running_loop()
+                if loop.is_running():
+                    raise RuntimeError("Cannot call synchronous `estimate_cost` inside loop.")
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+            
+            # Retrieve mock sources
+            sources = loop.run_until_complete(self.retriever.retrieve(query))
+            context_text = "\n".join([s.text for s in sources])
+            
+            estimate = estimator.estimate_generation(
+                prompt=f"{context_text}\n{query}",
+                approx_output_tokens=500
+            )
+            return estimate.__str__()
+
+
+
+__all__ = ["RAGBox"]

ragbox/cli.py

@@ -0,0 +1,106 @@
+import typer
+import asyncio
+from pathlib import Path
+from rich.console import Console
+from rich.panel import Panel
+from rich.markdown import Markdown
+
+from ragbox import RAGBox
+from ragbox.config.defaults import Settings
+
+app = typer.Typer(
+    name="ragbox",
+    help="RAGBox: Zero-Configuration Self-Building Agentic RAG System",
+    add_completion=False,
+)
+
+console = Console()
+
+@app.command()
+def init(
+    document_dir: Path = typer.Argument(
+        ...,
+        exists=True,
+        file_okay=False,
+        dir_okay=True,
+        help="Directory containing local documents to index.",
+    ),
+    config_file: Path = typer.Option(
+        None,
+        "--config",
+        "-c",
+        help="Path to an optional JSON/YAML configuration file.",
+    )
+):
+    """
+    Initialize a RAGBox instance in the specified directory.
+    
+    This command blocks until the initial index and knowledge graph are built.
+    Afterward, it starts the background daemon to watch for changes.
+    """
+    console.print(Panel(f"Initializing RAGBox at [cyan]{document_dir}[/cyan]", border_style="green"))
+    
+    # In a real implementation this might parse config_file into Settings
+    settings = Settings()
+    
+    try:
+        # Pre-flight cost estimation
+        rg_pre = RAGBox(document_dir, config=settings)
+        cost_estimate = rg_pre.estimate_cost(query=None)
+        console.print(Panel(f"[bold yellow]{cost_estimate}[/bold yellow]", title="Cost Estimate"))
+        
+        # RAGBox __init__ calls _ensure_built which builds synchronously implicitly via asyncio
+        with console.status("[bold green]Building Vector Store and Knowledge Graph...", spinner="dots"):
+            rg = rg_pre # It's already built by the init above
+        console.print("[bold green]✓ Indexing complete.[/bold green] RAGBox is ready for queries.")
+        console.print("The self-healing daemon is now active in the background for this session.")
+        
+        # Keep process alive to allow watchdog to run, or exit if init is purely 
+        # meant to pre-compute. Because RAGBox design spins up daemon threads, 
+        # we'll keep it alive here if we want continuous syncing.
+        console.print("[yellow]Press Ctrl+C to stop the daemon.[/yellow]")
+        try:
+            while True:
+                # Wait forever
+                asyncio.run(asyncio.sleep(3600))
+        except KeyboardInterrupt:
+            console.print("Daemon stopped.")
+        
+    except Exception as e:
+        console.print(f"[bold red]Error initializing RAGBox: {e}[/bold red]")
+        raise typer.Exit(code=1)
+
+@app.command()
+def query(
+    question: str = typer.Argument(..., help="The question to ask the RAG system."),
+    document_dir: Path = typer.Option(
+        ".",
+        "--dir",
+        "-d",
+        help="Directory where RAGBox is initialized.",
+        exists=True,
+        dir_okay=True,
+    )
+):
+    """
+    Query an existing RAGBox instance.
+    """
+    console.print(f"[bold blue]Querying RAGBox in {document_dir}...[/bold blue]\n")
+    try:
+        rg = RAGBox(document_dir)
+        
+        # Estimation
+        cost_estimate = rg.estimate_cost(query=question)
+        console.print(f"[yellow]Estimation:[/yellow] {cost_estimate}")
+        
+        with console.status("[bold green]Agentic Orchestrator thinking...", spinner="dots"):
+            answer = rg.query(question)
+            
+        console.print(Panel(Markdown(answer), title="Answer", border_style="blue"))
+        
+    except Exception as e:
+        console.print(f"[bold red]Error querying RAGBox: {e}[/bold red]")
+        raise typer.Exit(code=1)
+
+if __name__ == "__main__":
+    app()

ragbox/config/__init__.py

@@ -0,0 +1,3 @@
+"""Config module initialization."""
+from .defaults import Settings
+__all__ = ["Settings"]

ragbox/config/defaults.py

@@ -0,0 +1,28 @@
+"""
+Zero-Configuration Settings via Pydantic using dotenv.
+"""
+from typing import Optional
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic import Field
+
+class Settings(BaseSettings):
+    """
+    RAGBox Configuration.
+    Prioritizes Env Vars, falls back to .env if present, otherwise defaults.
+    """
+    model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8', extra='ignore')
+
+    # API Keys
+    openai_api_key: Optional[str] = Field(default=None, description="OpenAI API Key")
+    anthropic_api_key: Optional[str] = Field(default=None, description="Anthropic API Key")
+    groq_api_key: Optional[str] = Field(default=None, description="Groq API Key")
+    pinecone_api_key: Optional[str] = Field(default=None, description="Pinecone API Key")
+
+    # Local paths
+    local_model_path: str = Field(default="models/llama-3.1-8b-instruct.gguf", description="Path for local LLM")
+    chroma_db_dir: str = Field(default="./chroma_db", description="Path for local vector DB")
+    
+    # Behaviors
+    log_level: str = "INFO"
+    max_agent_steps: int = 5
+    graph_community_level: int = 1

ragbox/core/__init__.py

@@ -0,0 +1,18 @@
+"""
+Core RAGBox mechanisms.
+"""
+from .agentic_orchestrator import AgenticOrchestrator
+from .chunking_engine import ChunkingEngine
+from .document_processor import DocumentProcessorRouter
+from .knowledge_graph import OptimizedKnowledgeGraph
+from .retrieval_fusion import RetrievalFusionEngine
+from .self_healing import SelfHealer
+
+__all__ = [
+    "AgenticOrchestrator",
+    "ChunkingEngine",
+    "DocumentProcessorRouter",
+    "OptimizedKnowledgeGraph",
+    "RetrievalFusionEngine",
+    "SelfHealer",
+]

ragbox/core/agentic_orchestrator.py

@@ -0,0 +1,178 @@
+"""
+Layer 6: AGENTIC ORCHESTRATION
+Query classification, dynamic routing, ReAct planning, and synthesis.
+"""
+import time
+from typing import Any, Optional
+from loguru import logger
+
+from ragbox.models.queries import Answer, RAGStrategy
+from ragbox.core.retrieval_fusion import RetrievalFusionEngine
+from ragbox.utils.llm_clients import LLMClient
+
+class AgenticOrchestrator:
+    def __init__(self, 
+                 retriever: RetrievalFusionEngine, 
+                 llm_client: LLMClient,
+                 document_processor: Any,
+                 knowledge_graph: Any):
+        self.retriever = retriever
+        self.llm = llm_client
+        self.doc_processor = document_processor
+        self.kg = knowledge_graph
+
+    async def _classify_query(self, query: str) -> RAGStrategy:
+        """Classify query to route to correct strategy."""
+        schema = {
+            "strategy": "vector | agentic | graph | multi_query",
+            "reasoning": "string"
+        }
+        prompt = f"""
+        Classify this query into one of the following retrieval strategies based on intent:
+        - 'vector': General semantic search, asking "What is X?"
+        - 'graph': Relationships, connections ("How does X relate to Y?")
+        - 'multi_query': Broad comparisons, multi-step ("Compare A and B", "How do I deploy?")
+        - 'agentic': Complex reasoning requiring outside tool usage or sequence of thoughts.
+
+        Query: {query}
+        """
+        
+        try:
+            res = await self.llm.agenerate_structured(prompt, schema)
+            strategy_str = res.get("strategy", "vector").lower()
+            if strategy_str == "graph":
+                return RAGStrategy.GRAPH
+            elif strategy_str == "agentic":
+                return RAGStrategy.AGENTIC
+            elif strategy_str == "multi_query":
+                return RAGStrategy.MULTI_QUERY
+            return RAGStrategy.VECTOR
+        except Exception as e:
+            logger.warning(f"Classification failed: {e}. Defaulting to VECTOR.")
+            return RAGStrategy.VECTOR
+
+    async def _expand_query(self, query: str) -> list[str]:
+        """Expand one query into multiple semantic variations."""
+        prompt = f"""
+        Generate 3 distinct semantic variations of this query to improve search recall. 
+        Return ONLY the queries separated by newlines.
+        Query: {query}
+        """
+        try:
+            res = await self.llm.agenerate(prompt, system="You are an expert search query generator.")
+            queries = [q.strip("- \t1234567890.") for q in res.split('\n') if q.strip()]
+            queries = [q for q in queries if q]
+            if not queries:
+                 return [query]
+            return queries[:3]
+        except Exception as e:
+            logger.warning(f"Multi-query expansion failed: {e}")
+            return [query]
+
+    async def execute(self, query_text: str, force_strategy: Optional[RAGStrategy] = None) -> Answer:
+        """Execute the end-to-end RAG pipeline."""
+        start_time = time.time()
+        
+        strategy = force_strategy or await self._classify_query(query_text)
+        logger.info(f"Executing query '{query_text}' with strategy {strategy.name}")
+        
+        if strategy == RAGStrategy.AGENTIC:
+            answer = await self._execute_agentic(query_text)
+        elif strategy == RAGStrategy.GRAPH:
+            answer = await self._execute_graph(query_text)
+        elif strategy == RAGStrategy.MULTI_QUERY:
+            answer = await self._execute_multi_query(query_text)
+        else:
+            answer = await self._execute_vector(query_text)
+            
+        answer.execution_time_ms = (time.time() - start_time) * 1000
+        return answer
+
+    async def _execute_vector(self, query: str) -> Answer:
+        sources = await self.retriever.retrieve(query)
+        context = "\n\n---\n\n".join([s.text for s in sources])
+        
+        prompt = f"Answer the query based ONLY on the following context.\n\nContext:\n{context}\n\nQuery: {query}"
+        response = await self.llm.agenerate(prompt, system="You are an expert Q&A engine. Be concise and accurate.")
+        
+        return Answer(
+            query=query,
+            content=response,
+            sources=sources,
+            strategy_used=RAGStrategy.VECTOR
+        )
+
+    async def _execute_multi_query(self, query: str) -> Answer:
+        expanded_queries = await self._expand_query(query)
+        logger.info(f"Expanded initial query into: {expanded_queries}")
+        
+        all_sources = []
+        for q in [query] + expanded_queries:
+             q_sources = await self.retriever.retrieve(q, top_k=3)
+             all_sources.extend(q_sources)
+             
+        # Deduplicate sources by text chunk ID or content to avoid noise
+        seen = set()
+        unique_sources = []
+        for s in all_sources:
+             # simple dedup
+             if s.text not in seen:
+                  seen.add(s.text)
+                  unique_sources.append(s)
+                  
+        # Limit to top N unique (since reranker will do the heavy lifting later in the pipeline anyway)
+        unique_sources = unique_sources[:10]
+        
+        context = "\n\n---\n\n".join([s.text for s in unique_sources])
+        prompt = f"Answer the query comprehensively based ONLY on the following context derived from multi-query expansion.\n\nContext:\n{context}\n\nQuery: {query}"
+        response = await self.llm.agenerate(prompt, system="You are an expert Q&A engine analyzing across multiple perspectives.")
+        
+        return Answer(
+            query=query,
+            content=response,
+            sources=unique_sources,
+            strategy_used=RAGStrategy.MULTI_QUERY
+        )
+
+
+    async def _execute_graph(self, query: str) -> Answer:
+        sources = await self.retriever.retrieve(query)
+        context = "\n\n---\n\n".join([s.text for s in sources])
+        
+        prompt = f"Answer this query using the graph summaries and retrieved text provided below:\n\nContext:\n{context}\n\nQuery: {query}"
+        response = await self.llm.agenerate(prompt, system="You are a graph-aware reasoning agent.")
+        
+        return Answer(
+            query=query,
+            content=response,
+            sources=sources,
+            strategy_used=RAGStrategy.GRAPH
+        )
+
+    async def _execute_agentic(self, query: str) -> Answer:
+        history = []
+        max_steps = 3
+        sources = []
+        
+        for step in range(max_steps):
+            prompt = f"Query: {query}\nHistory: {history}\n\nAction (SEARCH <term> or ANSWER <final answer>):"
+            action = await self.llm.agenerate(prompt, system="You are a ReAct agent. You can SEARCH or ANSWER.")
+            
+            if "ANSWER" in action:
+                final_ans = action.split("ANSWER")[-1].strip(': \n')
+                return Answer(
+                    query=query,
+                    content=final_ans,
+                    sources=sources,
+                    strategy_used=RAGStrategy.AGENTIC
+                )
+            elif "SEARCH" in action:
+                search_term = action.split("SEARCH")[-1].strip(': \n')
+                step_sources = await self.retriever.retrieve(search_term, top_k=2)
+                sources.extend(step_sources)
+                context = "\n".join([s.text for s in step_sources])
+                history.append(f"Searched: {search_term}\nFound context length: {len(context)}")
+            else:
+                break
+                
+        return await self._execute_vector(query)

ragbox/core/chunking_engine.py

@@ -0,0 +1,169 @@
+"""
+Layer 3: SELF-OPTIMIZING CHUNKING
+Configures and dispatches optimal chunking strategy.
+"""
+from abc import ABC, abstractmethod
+from typing import List, Any
+import hashlib
+from loguru import logger
+import asyncio
+
+from ragbox.models.documents import Document
+from ragbox.models.chunks import Chunk, TextChunk
+from ragbox.utils.llm_clients import LLMClient
+from ragbox.utils.embeddings import EmbeddingProvider
+
+class ChunkingStrategy(ABC):
+    @abstractmethod
+    def chunk(self, document: Document) -> List[Chunk]:
+        """Synchronously chunk a document."""
+        pass
+
+
+class FixedChunker(ChunkingStrategy):
+    """Fallback basic text chunker."""
+    def __init__(self, chunk_size: int = 1000, overlap: int = 200):
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    def chunk(self, document: Document) -> List[Chunk]:
+        chunks = []
+        text = document.content
+        if not text:
+            return chunks
+        
+        start = 0
+        while start < len(text):
+            end = start + self.chunk_size
+            chunk_text = text[start:end]
+            chunk_id = hashlib.sha256(f"{document.id}_{start}".encode()).hexdigest()
+            chunks.append(TextChunk(
+                id=chunk_id,
+                document_id=document.id,
+                content=chunk_text,
+                metadata={"start_idx": start, "end_idx": end, "strategy": "fixed"}
+            ))
+            start += self.chunk_size - self.overlap
+        return chunks
+
+
+class SentenceChunker(ChunkingStrategy):
+    """Chunker that attempts to split on sentences to preserve context."""
+    def __init__(self, max_sentences: int = 5):
+        self.max_sentences = max_sentences
+
+    def chunk(self, document: Document) -> List[Chunk]:
+        chunks = []
+        text = document.content
+        if not text:
+            return chunks
+            
+        import re
+        # very naive sentence splitting
+        sentences = re.split(r'(?<=[.!?]) +', text)
+        
+        current_chunk = []
+        for i, sentence in enumerate(sentences):
+            current_chunk.append(sentence)
+            if len(current_chunk) >= self.max_sentences or i == len(sentences) - 1:
+                chunk_text = " ".join(current_chunk)
+                chunk_id = hashlib.sha256(f"{document.id}_{i}".encode()).hexdigest()
+                chunks.append(TextChunk(
+                    id=chunk_id,
+                    document_id=document.id,
+                    content=chunk_text,
+                    metadata={"strategy": "sentence"}
+                ))
+                # Slight overlap: keep the last sentence for the next chunk
+                current_chunk = [current_chunk[-1]] if len(current_chunk) > 0 else []
+                
+        return chunks
+
+
+class SelfOptimizingChunker:
+    """Evaluates and selects best strategy."""
+    def __init__(self, llm_client: LLMClient):
+        self.llm = llm_client
+        self.strategies = {
+            "fixed_small": FixedChunker(chunk_size=500, overlap=100),
+            "fixed_large": FixedChunker(chunk_size=1500, overlap=300),
+            "sentence": SentenceChunker(max_sentences=8)
+        }
+
+    async def optimize(self, documents: List[Document]) -> ChunkingStrategy:
+        """Auto-evaluates which strategy is best by sampling."""
+        if not documents:
+             return self.strategies["fixed_large"]
+             
+        # Sample document 
+        sample_doc = documents[0]
+        # Only sample first 5000 chars to save tokens
+        sample_doc_truncated = Document(id="sample", path=sample_doc.path, content=sample_doc.content[:5000])
+        
+        logger.debug(f"Auto-optimizing chunking strategy based on {sample_doc.path.name}...")
+        
+        # Test strategies
+        evaluations = {}
+        for name, strategy in self.strategies.items():
+            chunks = strategy.chunk(sample_doc_truncated)
+            if not chunks:
+                 continue
+                 
+            # Take the first ~3 chunks to show the LLM
+            sample_chunks_text = "\n---\n".join([c.content for c in chunks[:3]])
+            
+            prompt = f"""
+            You are evaluating text chunking strategies for a RAG system.
+            Given this sample of document chunks extracted using strategy '{name}', 
+            score it from 1 to 10 on how semantically coherent and self-contained the chunks are.
+            Return ONLY the integer score.
+            
+            Chunks:
+            {sample_chunks_text}
+            """
+            
+            try:
+                score_str = await self.llm.agenerate(prompt, system="You return only a single integer between 1 and 10.")
+                import re
+                numbers = re.findall(r'\d+', score_str)
+                score = int(numbers[0]) if numbers else 5
+                evaluations[name] = score
+            except Exception as e:
+                logger.warning(f"Chunk evaluation failed for {name}: {e}")
+                evaluations[name] = 5
+                
+        if not evaluations:
+             return self.strategies["fixed_large"]
+             
+        best_strategy_name = max(evaluations, key=evaluations.get)
+        logger.info(f"Auto-Optimized Chunking: Selected '{best_strategy_name}' (Scores: {evaluations})")
+        return self.strategies[best_strategy_name]
+
+
+class ChunkingEngine:
+    """Entry point for Layer 3."""
+    def __init__(self, llm_client: LLMClient, embedding_provider: EmbeddingProvider):
+        self.llm = llm_client
+        self.embedding_provider = embedding_provider
+        self.optimizer = SelfOptimizingChunker(self.llm)
+
+    async def chunk(self, document: Document) -> List[Chunk]:
+        """Apply optimal chunking to document and embed chunks."""
+        # Use optimizer to pick the best strategy for this specific document
+        strategy = await self.optimizer.optimize([document])
+        chunks = await asyncio.to_thread(strategy.chunk, document)
+        
+        if chunks:
+            try:
+                # Use Late Chunking / Contextual Retrieval
+                embeddings = await self.embedding_provider.embed_chunks_with_context(
+                    document=document, 
+                    chunks=chunks, 
+                    llm_client=self.llm
+                )
+                for chunk, emb in zip(chunks, embeddings):
+                    chunk.metadata["embedding"] = emb
+            except Exception as e:
+                logger.error(f"Failed to embed chunks for document {document.path}: {e}")
+                
+        return chunks