Merge branch 'worktree-agent-a36bc4ac'

Author: shantanu patil

api/context_budget.py

@@ -0,0 +1,114 @@
+"""Context Budget Manager --- manages token allocation for LLM context assembly."""
+
+import logging
+from typing import List, Dict, Optional, Any, Callable
+
+logger = logging.getLogger(__name__)
+
+# Known context window sizes by provider and model
+CONTEXT_WINDOWS = {
+    "google": {
+        "gemini-2.5-flash": 1_048_576,
+        "gemini-2.5-flash-lite": 1_048_576,
+        "gemini-2.5-pro": 1_048_576,
+        "gemini-2.0-flash": 1_048_576,
+        "gemini-1.5-pro": 2_097_152,
+        "gemini-1.5-flash": 1_048_576,
+    },
+    "openai": {
+        "gpt-4o": 128_000,
+        "gpt-4o-mini": 128_000,
+        "gpt-4.1": 1_000_000,
+        "gpt-4.1-mini": 1_000_000,
+        "o3-mini": 200_000,
+    },
+    "openrouter": {},  # varies by model
+    "ollama": {},  # varies by model, default 32K
+    "bedrock": {},
+    "azure": {},
+    "dashscope": {},
+}
+
+DEFAULT_CONTEXT_WINDOW = 128_000
+DEFAULT_OUTPUT_RESERVE = 8_192
+
+
+class ContextBudgetManager:
+    """Manages token budget for LLM context assembly."""
+
+    def get_context_window(self, provider: str, model: str) -> int:
+        """Get the context window size for a provider/model combo."""
+        provider_windows = CONTEXT_WINDOWS.get(provider, {})
+        # Try exact match first
+        if model in provider_windows:
+            return provider_windows[model]
+        # Try prefix match (e.g., "gemini-2.5-flash-preview" matches "gemini-2.5-flash")
+        for known_model, window in provider_windows.items():
+            if model.startswith(known_model):
+                return window
+        return DEFAULT_CONTEXT_WINDOW
+
+    def get_context_budget(self, provider: str, model: str,
+                           prompt_tokens: int,
+                           output_reserve: int = DEFAULT_OUTPUT_RESERVE) -> int:
+        """Calculate available tokens for RAG context."""
+        window = self.get_context_window(provider, model)
+        budget = window - prompt_tokens - output_reserve
+        logger.info(f"Context budget: window={window}, prompt={prompt_tokens}, "
+                    f"reserve={output_reserve}, available={budget}")
+        return max(budget, 0)
+
+    def get_dynamic_top_k(self, provider: str, model: str,
+                          avg_chunk_tokens: int = 600) -> int:
+        """Calculate how many chunks to retrieve based on context window."""
+        window = self.get_context_window(provider, model)
+        # Use at most 30% of context window for retrieved chunks
+        chunk_budget = int(window * 0.3)
+        top_k = max(20, min(200, chunk_budget // avg_chunk_tokens))
+        logger.info(f"Dynamic top_k for {provider}/{model}: {top_k}")
+        return top_k
+
+    def assemble_context(self, documents: List[Dict],
+                         budget_tokens: int,
+                         count_tokens_fn: Callable[[str], int]) -> str:
+        """Greedily pack documents into the token budget, highest relevance first.
+
+        Args:
+            documents: List of dicts with 'content' and optionally 'score', 'file_path'
+            budget_tokens: Maximum tokens for context
+            count_tokens_fn: Function that counts tokens in a string
+
+        Returns:
+            Assembled context string fitting within budget
+        """
+        # Sort by relevance score (highest first)
+        sorted_docs = sorted(documents, key=lambda d: d.get('score', 0), reverse=True)
+
+        assembled = []
+        used_tokens = 0
+
+        for doc in sorted_docs:
+            content = doc.get('content', '')
+            doc_tokens = count_tokens_fn(content)
+
+            if used_tokens + doc_tokens > budget_tokens:
+                # Try to fit a truncated version
+                remaining = budget_tokens - used_tokens
+                if remaining > 100:  # Only include if meaningful content fits
+                    # Rough truncation by ratio
+                    ratio = remaining / doc_tokens
+                    truncated = content[:int(len(content) * ratio)]
+                    assembled.append(truncated + "\n... [truncated]")
+                break
+
+            file_path = doc.get('file_path', 'unknown')
+            assembled.append(f"--- {file_path} ---\n{content}")
+            used_tokens += doc_tokens
+
+        logger.info(f"Assembled context: {len(assembled)} docs, ~{used_tokens} tokens "
+                    f"(budget: {budget_tokens})")
+        return "\n\n".join(assembled)
+
+
+# Module-level singleton
+context_budget_manager = ContextBudgetManager()

api/rag.py

@@ -8,6 +8,7 @@
 import adalflow as adal
 
 from api.tools.embedder import get_embedder
+from api.context_budget import context_budget_manager
 from api.prompts import RAG_SYSTEM_PROMPT as system_prompt, RAG_TEMPLATE
 
 # Create our own implementation of the conversation classes
@@ -382,13 +383,20 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_
         try:
             # Use the appropriate embedder for retrieval
             retrieve_embedder = self.query_embedder if self.is_ollama_embedder else self.embedder
+
+            # Calculate dynamic top_k based on model context window
+            retriever_kwargs = dict(configs["retriever"])
+            dynamic_top_k = self._get_dynamic_top_k()
+            if dynamic_top_k is not None:
+                retriever_kwargs["top_k"] = dynamic_top_k
+
             self.retriever = FAISSRetriever(
-                **configs["retriever"],
+                **retriever_kwargs,
                 embedder=retrieve_embedder,
                 documents=self.transformed_docs,
                 document_map_func=lambda doc: doc.vector,
             )
-            logger.info("FAISS retriever created successfully")
+            logger.info(f"FAISS retriever created successfully (top_k={retriever_kwargs.get('top_k', 'default')})")
         except Exception as e:
             logger.error(f"Error creating FAISS retriever: {str(e)}")
             # Try to provide more specific error information
@@ -413,6 +421,27 @@ def prepare_retriever(self, repo_url_or_path: str, type: str = "github", access_
                 logger.error(f"Sample embedding sizes: {', '.join(sizes)}")
             raise
 
+    def _get_dynamic_top_k(self) -> int:
+        """Calculate dynamic top_k based on the model's context window.
+
+        Uses the context_budget_manager to determine how many chunks to
+        retrieve based on the provider and model. Returns None if the
+        provider/model are not set (falls back to config default).
+        """
+        if self.provider and self.model:
+            try:
+                # Get average chunk size from text_splitter config
+                chunk_size = configs.get("text_splitter", {}).get("chunk_size", 350)
+                # Rough estimate: 1 word ~ 1.3 tokens, so chunk_tokens ~ chunk_size * 1.3
+                avg_chunk_tokens = int(chunk_size * 1.3)
+                dynamic_top_k = context_budget_manager.get_dynamic_top_k(
+                    self.provider, self.model, avg_chunk_tokens=avg_chunk_tokens
+                )
+                return dynamic_top_k
+            except Exception as e:
+                logger.warning(f"Failed to calculate dynamic top_k: {e}")
+        return None
+
     def call(self, query: str, language: str = "en") -> Tuple[List]:
         """
         Process a query using RAG.

api/rag_session.py

@@ -0,0 +1,70 @@
+"""RAG Session Manager --- caches RAG instances for reuse across page generations."""
+
+import time
+import threading
+import logging
+from typing import Optional, Dict, Tuple, Any
+
+logger = logging.getLogger(__name__)
+
+
+class RAGSessionManager:
+    """Caches RAG instances keyed by repo URL to avoid rebuilding FAISS index per page."""
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._sessions: Dict[str, Tuple[Any, float]] = {}
+                    cls._instance._ttl = 3600  # 1 hour TTL
+                    cls._instance._max_sessions = 10
+        return cls._instance
+
+    def get_session_key(self, repo_url: str, embedder_type: str = "default") -> str:
+        """Generate a cache key for a RAG session."""
+        return f"{repo_url}:{embedder_type}"
+
+    def get(self, key: str):
+        """Get a cached RAG instance if it exists and hasn't expired."""
+        with self._lock:
+            if key in self._sessions:
+                rag, last_access = self._sessions[key]
+                if time.time() - last_access < self._ttl:
+                    self._sessions[key] = (rag, time.time())
+                    logger.info(f"RAG session cache hit for {key}")
+                    return rag
+                else:
+                    # Expired
+                    del self._sessions[key]
+                    logger.info(f"RAG session expired for {key}")
+        return None
+
+    def put(self, key: str, rag_instance):
+        """Cache a RAG instance."""
+        with self._lock:
+            # Evict oldest if at capacity
+            if len(self._sessions) >= self._max_sessions and key not in self._sessions:
+                oldest_key = min(self._sessions, key=lambda k: self._sessions[k][1])
+                del self._sessions[oldest_key]
+                logger.info(f"Evicted oldest RAG session: {oldest_key}")
+
+            self._sessions[key] = (rag_instance, time.time())
+            logger.info(f"Cached RAG session for {key}")
+
+    def invalidate(self, key: str):
+        """Remove a cached session."""
+        with self._lock:
+            self._sessions.pop(key, None)
+
+    def clear(self):
+        """Clear all cached sessions."""
+        with self._lock:
+            self._sessions.clear()
+
+
+# Module-level singleton
+rag_session_manager = RAGSessionManager()

api/websocket_wiki.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import os
 from typing import List, Optional, Dict, Any
@@ -23,7 +24,8 @@
 from api.openrouter_client import OpenRouterClient
 from api.azureai_client import AzureAIClient
 from api.dashscope_client import DashscopeClient
-from api.rag import RAG
+from api.rag import RAG, Memory
+from api.rag_session import rag_session_manager
 
 # Configure logging
 from api.logging_config import setup_logging
@@ -60,6 +62,70 @@ class ChatCompletionRequest(BaseModel):
     included_dirs: Optional[str] = Field(None, description="Comma-separated list of directories to include exclusively")
     included_files: Optional[str] = Field(None, description="Comma-separated list of file patterns to include exclusively")
 
+async def generate_with_retry(rag, query, context_docs, provider, model, language="en", max_retries=3):
+    """Generate content with retry and context reduction on failure.
+
+    On token limit errors, reduces context by 50% per retry.
+    On transient errors (timeout, 503, 429), retries with exponential backoff.
+    Non-retryable errors are raised immediately.
+
+    Args:
+        rag: RAG instance to use for generation
+        query: The user query
+        context_docs: List of retrieved documents
+        provider: AI provider name
+        model: Model name
+        language: Language code for content generation
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        Retrieved documents result from RAG
+    """
+    context_fraction = 1.0
+
+    for attempt in range(max_retries):
+        try:
+            # Reduce context on retries
+            if context_fraction < 1.0 and context_docs:
+                reduced_count = max(1, int(len(context_docs) * context_fraction))
+                docs_to_use = context_docs[:reduced_count]
+                logger.info(f"Using {len(docs_to_use)}/{len(context_docs)} context docs "
+                           f"({context_fraction:.0%})")
+            else:
+                docs_to_use = context_docs
+
+            result = rag(query, language=language)
+            return result
+
+        except Exception as e:
+            error_str = str(e).lower()
+
+            # Token limit errors -- reduce context
+            if any(phrase in error_str for phrase in [
+                'maximum context length', 'token limit', 'too many tokens',
+                'content too large', 'request too large', 'input too long'
+            ]):
+                context_fraction *= 0.5
+                logger.warning(f"Token limit hit, reducing context to {context_fraction:.0%} "
+                             f"(attempt {attempt + 1}/{max_retries})")
+                continue
+
+            # Transient errors -- retry with backoff
+            if any(phrase in error_str for phrase in [
+                'timeout', 'connection', '503', '502', '429', 'rate limit'
+            ]):
+                wait_time = (2 ** attempt)  # 1s, 2s, 4s
+                logger.warning(f"Transient error, retrying in {wait_time}s "
+                             f"(attempt {attempt + 1}/{max_retries}): {e}")
+                await asyncio.sleep(wait_time)
+                continue
+
+            # Non-retryable error
+            raise
+
+    raise Exception(f"Failed after {max_retries} retries with context at {context_fraction:.0%}")
+
+
 async def handle_websocket_chat(websocket: WebSocket):
     """
     Handle WebSocket connection for chat completions.
@@ -83,10 +149,8 @@ async def handle_websocket_chat(websocket: WebSocket):
                     logger.warning(f"Request exceeds recommended token limit ({tokens} > 7500)")
                     input_too_large = True
 
-        # Create a new RAG instance for this request
+        # Create or reuse a cached RAG instance for this request
         try:
-            request_rag = RAG(provider=request.provider, model=request.model)
-
             # Extract custom file filter parameters if provided
             excluded_dirs = None
             excluded_files = None
@@ -106,8 +170,28 @@ async def handle_websocket_chat(websocket: WebSocket):
                 included_files = [unquote(file_pattern) for file_pattern in request.included_files.split('\n') if file_pattern.strip()]
                 logger.info(f"Using custom included files: {included_files}")
 
-            request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files)
-            logger.info(f"Retriever prepared for {request.repo_url}")
+            # Check for a cached RAG session (only when no custom file filters)
+            has_custom_filters = any([excluded_dirs, excluded_files, included_dirs, included_files])
+            from api.config import get_embedder_type
+            embedder_type = get_embedder_type()
+            session_key = rag_session_manager.get_session_key(request.repo_url, embedder_type) if not has_custom_filters else None
+            request_rag = rag_session_manager.get(session_key) if session_key else None
+
+            if request_rag is not None:
+                # Reuse cached RAG instance, update provider/model for this request
+                request_rag.provider = request.provider
+                request_rag.model = request.model
+                # Reset memory for this new conversation
+                request_rag.memory = Memory()
+                logger.info(f"Reusing cached RAG session for {request.repo_url}")
+            else:
+                # Create a new RAG instance
+                request_rag = RAG(provider=request.provider, model=request.model)
+                request_rag.prepare_retriever(request.repo_url, request.type, request.token, excluded_dirs, excluded_files, included_dirs, included_files)
+                # Cache the session if no custom filters were used
+                if session_key:
+                    rag_session_manager.put(session_key, request_rag)
+                logger.info(f"Created new RAG session for {request.repo_url}")
         except ValueError as e:
             if "No valid documents with embeddings found" in str(e):
                 logger.error(f"No valid embeddings found: {str(e)}")
@@ -202,10 +286,14 @@ async def handle_websocket_chat(websocket: WebSocket):
                     rag_query = f"Contexts related to {request.filePath}"
                     logger.info(f"Modified RAG query to focus on file: {request.filePath}")
 
-                # Try to perform RAG retrieval
+                # Try to perform RAG retrieval with retry logic
                 try:
-                    # This will use the actual RAG implementation
-                    retrieved_documents = request_rag(rag_query, language=request.language)
+                    # Use retry wrapper for resilient retrieval
+                    retrieved_documents = await generate_with_retry(
+                        request_rag, rag_query, None,
+                        request.provider, request.model,
+                        language=request.language
+                    )
 
                     if retrieved_documents and retrieved_documents[0].documents:
                         # Format context for the prompt in a more structured way