agent-brain-rag 1.1.0__py3-none-any.whl

doc_serve_server/indexing/embedding.py

@@ -0,0 +1,274 @@
+"""Embedding generation using OpenAI's text-embedding models."""
+
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Optional
+
+from anthropic import AsyncAnthropic
+from openai import AsyncOpenAI
+
+from doc_serve_server.config import settings
+
+from .chunking import TextChunk
+
+logger = logging.getLogger(__name__)
+
+
+class EmbeddingGenerator:
+    """
+    Generates embeddings using OpenAI's embedding models.
+
+    Supports batch processing with configurable batch sizes
+    and automatic retry on rate limits.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        batch_size: Optional[int] = None,
+    ):
+        """
+        Initialize the embedding generator.
+
+        Args:
+            api_key: OpenAI API key. Defaults to config value.
+            model: Embedding model name. Defaults to config value.
+            batch_size: Number of texts to embed per API call. Defaults to 100.
+        """
+        self.model = model or settings.EMBEDDING_MODEL
+        self.batch_size = batch_size or settings.EMBEDDING_BATCH_SIZE
+
+        # Initialize OpenAI async client
+        self.client = AsyncOpenAI(
+            api_key=api_key or settings.OPENAI_API_KEY,
+        )
+
+        # Initialize Anthropic client for summarization
+        self.anthropic_client = AsyncAnthropic(
+            api_key=settings.ANTHROPIC_API_KEY,
+        )
+
+        # Initialize prompt template
+        self.summary_prompt_template = (
+            "You are an expert software engineer analyzing source code. "
+            "Provide a concise 1-2 sentence summary of what this code does. "
+            "Focus on the functionality, purpose, and behavior. "
+            "Be specific about inputs, outputs, and side effects. "
+            "Ignore implementation details and focus on what the code accomplishes.\n\n"
+            "Code to summarize:\n{context_str}\n\n"
+            "Summary:"
+        )
+
+    async def embed_text(self, text: str) -> list[float]:
+        """
+        Generate embedding for a single text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding vector as list of floats.
+        """
+        response = await self.client.embeddings.create(
+            model=self.model,
+            input=text,
+        )
+        return response.data[0].embedding
+
+    async def embed_texts(
+        self,
+        texts: list[str],
+        progress_callback: Optional[Callable[[int, int], Awaitable[None]]] = None,
+    ) -> list[list[float]]:
+        """
+        Generate embeddings for multiple texts.
+
+        Args:
+            texts: List of texts to embed.
+            progress_callback: Optional callback(processed, total) for progress.
+
+        Returns:
+            List of embedding vectors.
+        """
+        if not texts:
+            return []
+
+        all_embeddings: list[list[float]] = []
+
+        # Process in batches to respect API limits
+        for i in range(0, len(texts), self.batch_size):
+            batch = texts[i : i + self.batch_size]
+
+            try:
+                response = await self.client.embeddings.create(
+                    model=self.model,
+                    input=batch,
+                )
+
+                # Extract embeddings in order
+                batch_embeddings = [item.embedding for item in response.data]
+                all_embeddings.extend(batch_embeddings)
+
+                if progress_callback:
+                    await progress_callback(
+                        min(i + self.batch_size, len(texts)),
+                        len(texts),
+                    )
+
+                logger.debug(
+                    f"Generated embeddings for batch {i // self.batch_size + 1} "
+                    f"({len(batch)} texts)"
+                )
+
+            except Exception as e:
+                logger.error(f"Failed to generate embeddings for batch: {e}")
+                raise
+
+        return all_embeddings
+
+    async def embed_chunks(
+        self,
+        chunks: list[TextChunk],
+        progress_callback: Optional[Callable[[int, int], Awaitable[None]]] = None,
+    ) -> list[list[float]]:
+        """
+        Generate embeddings for a list of text chunks.
+
+        Args:
+            chunks: List of TextChunk objects.
+            progress_callback: Optional callback for progress updates.
+
+        Returns:
+            List of embedding vectors corresponding to each chunk.
+        """
+        texts = [chunk.text for chunk in chunks]
+        return await self.embed_texts(texts, progress_callback)
+
+    async def embed_query(self, query: str) -> list[float]:
+        """
+        Generate embedding for a search query.
+
+        This is a convenience wrapper around embed_text for queries.
+
+        Args:
+            query: The search query text.
+
+        Returns:
+            Query embedding vector.
+        """
+        return await self.embed_text(query)
+
+    def get_embedding_dimensions(self) -> int:
+        """
+        Get the expected embedding dimensions for the current model.
+
+        Returns:
+            Number of dimensions in the embedding vector.
+        """
+        # Known dimensions for OpenAI models
+        model_dimensions = {
+            "text-embedding-3-large": 3072,
+            "text-embedding-3-small": 1536,
+            "text-embedding-ada-002": 1536,
+        }
+        return model_dimensions.get(self.model, settings.EMBEDDING_DIMENSIONS)
+
+    def _get_summary_prompt_template(self) -> str:
+        """
+        Get the prompt template for code summarization.
+
+        Returns:
+            Prompt template string.
+        """
+        template = (
+            "You are an expert software engineer analyzing source code. "
+            "Provide a concise 1-2 sentence summary of what this code does. "
+            "Focus on the functionality, purpose, and behavior. "
+            "Be specific about inputs, outputs, and side effects. "
+            "Ignore implementation details and focus on what the code accomplishes.\n\n"
+            "Code to summarize:\n{context_str}\n\n"
+            "Summary:"
+        )
+        return template
+
+    async def generate_summary(self, code_text: str) -> str:
+        """
+        Generate a natural language summary of code using Claude.
+
+        Args:
+            code_text: The source code to summarize.
+
+        Returns:
+            Natural language summary of the code's functionality.
+        """
+        try:
+            # Use Claude directly with custom prompt
+            prompt = self.summary_prompt_template.format(context_str=code_text)
+
+            response = await self.anthropic_client.messages.create(
+                model=settings.CLAUDE_MODEL,
+                max_tokens=300,
+                temperature=0.1,  # Low temperature for consistent summaries
+                messages=[{"role": "user", "content": prompt}],
+            )
+
+            # Extract text from Claude response
+            summary = response.content[0].text  # type: ignore
+
+            if summary and len(summary) > 10:  # Ensure we got a meaningful summary
+                return summary
+            else:
+                logger.warning("Claude returned empty or too short summary")
+                return self._extract_fallback_summary(code_text)
+
+        except Exception as e:
+            logger.error(f"Failed to generate code summary: {e}")
+            # Fallback: try to extract from docstrings/comments
+            return self._extract_fallback_summary(code_text)
+
+    def _extract_fallback_summary(self, code_text: str) -> str:
+        """
+        Extract summary from docstrings or comments as fallback.
+
+        Args:
+            code_text: Source code to analyze.
+
+        Returns:
+            Extracted summary or empty string.
+        """
+        import re
+
+        # Try to find Python docstrings
+        docstring_match = re.search(r'""".*?"""', code_text, re.DOTALL)
+        if docstring_match:
+            docstring = docstring_match.group(0)[3:-3]  # Remove leading/trailing """
+            if len(docstring) > 10:  # Only use if substantial
+                return docstring[:200] + "..." if len(docstring) > 200 else docstring
+
+        # Try to find function/class comments
+        comment_match = re.search(
+            r"#.*(?:function|class|method|def)", code_text, re.IGNORECASE
+        )
+        if comment_match:
+            return comment_match.group(0).strip("#").strip()
+
+        # Last resort: first line if it looks like a comment
+        lines = code_text.strip().split("\n")
+        first_line = lines[0].strip()
+        if first_line.startswith(("#", "//", "/*")):
+            return first_line.lstrip("#/*").strip()
+
+        return ""  # No summary available
+
+
+# Singleton instance
+_embedding_generator: Optional[EmbeddingGenerator] = None
+
+
+def get_embedding_generator() -> EmbeddingGenerator:
+    """Get the global embedding generator instance."""
+    global _embedding_generator
+    if _embedding_generator is None:
+        _embedding_generator = EmbeddingGenerator()
+    return _embedding_generator

doc_serve_server/locking.py

@@ -0,0 +1,133 @@
+"""File-based locking for doc-serve instances."""
+
+import fcntl
+import logging
+import os
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+LOCK_FILE = "doc-serve.lock"
+PID_FILE = "doc-serve.pid"
+
+# Module-level storage for lock file descriptors
+_lock_fds: dict[str, int] = {}
+
+
+def acquire_lock(state_dir: Path) -> bool:
+    """Acquire an exclusive lock for the state directory.
+
+    Non-blocking. Returns immediately if lock cannot be acquired.
+
+    Args:
+        state_dir: Path to the state directory.
+
+    Returns:
+        True if lock acquired, False if already held.
+    """
+    state_dir.mkdir(parents=True, exist_ok=True)
+    lock_path = state_dir / LOCK_FILE
+
+    try:
+        fd = os.open(str(lock_path), os.O_CREAT | os.O_WRONLY)
+        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+
+        # Write PID
+        pid_path = state_dir / PID_FILE
+        pid_path.write_text(str(os.getpid()))
+
+        # Store fd for later release
+        _lock_fds[str(state_dir)] = fd
+        logger.info(f"Lock acquired: {lock_path}")
+        return True
+
+    except OSError:
+        logger.warning(f"Lock already held: {lock_path}")
+        return False
+
+
+def release_lock(state_dir: Path) -> None:
+    """Release the lock for the state directory.
+
+    Args:
+        state_dir: Path to the state directory.
+    """
+    lock_path = state_dir / LOCK_FILE
+
+    fd = _lock_fds.pop(str(state_dir), None)
+    if fd is not None:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_UN)
+            os.close(fd)
+        except OSError:
+            pass
+
+    # Clean up files
+    for fname in [LOCK_FILE, PID_FILE]:
+        fpath = state_dir / fname
+        if fpath.exists():
+            try:
+                fpath.unlink()
+            except OSError:
+                pass
+
+    logger.info(f"Lock released: {lock_path}")
+
+
+def read_pid(state_dir: Path) -> Optional[int]:
+    """Read the PID from the PID file.
+
+    Args:
+        state_dir: Path to the state directory.
+
+    Returns:
+        PID value or None if file doesn't exist or is invalid.
+    """
+    pid_path = state_dir / PID_FILE
+    if not pid_path.exists():
+        return None
+    try:
+        return int(pid_path.read_text().strip())
+    except (ValueError, OSError):
+        return None
+
+
+def is_stale(state_dir: Path) -> bool:
+    """Check if the lock is stale (PID no longer alive).
+
+    Args:
+        state_dir: Path to the state directory.
+
+    Returns:
+        True if the lock is stale or no PID exists.
+    """
+    pid = read_pid(state_dir)
+    if pid is None:
+        return True
+    try:
+        os.kill(pid, 0)
+        return False  # Process is alive
+    except ProcessLookupError:
+        return True  # Process is dead
+    except PermissionError:
+        return False  # Process exists but we can't signal it
+
+
+def cleanup_stale(state_dir: Path) -> None:
+    """Clean up stale lock and PID files.
+
+    Only cleans up if the lock is determined to be stale.
+
+    Args:
+        state_dir: Path to the state directory.
+    """
+    if is_stale(state_dir):
+        for fname in [LOCK_FILE, PID_FILE, "runtime.json"]:
+            fpath = state_dir / fname
+            if fpath.exists():
+                try:
+                    fpath.unlink()
+                    logger.info(f"Cleaned stale file: {fpath}")
+                except OSError:
+                    pass

doc_serve_server/models/__init__.py

@@ -0,0 +1,18 @@
+"""Pydantic models for request/response handling."""
+
+from .health import HealthStatus, IndexingStatus
+from .index import IndexingState, IndexingStatusEnum, IndexRequest, IndexResponse
+from .query import QueryMode, QueryRequest, QueryResponse, QueryResult
+
+__all__ = [
+    "QueryMode",
+    "QueryRequest",
+    "QueryResponse",
+    "QueryResult",
+    "IndexRequest",
+    "IndexResponse",
+    "IndexingState",
+    "IndexingStatusEnum",
+    "HealthStatus",
+    "IndexingStatus",
+]

doc_serve_server/models/health.py

@@ -0,0 +1,126 @@
+"""Health status models."""
+
+from datetime import datetime, timezone
+from typing import Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class HealthStatus(BaseModel):
+    """Server health status response."""
+
+    status: Literal["healthy", "indexing", "degraded", "unhealthy"] = Field(
+        ...,
+        description="Current server health status",
+    )
+    message: Optional[str] = Field(
+        None,
+        description="Additional status message",
+    )
+    timestamp: datetime = Field(
+        default_factory=lambda: datetime.now(timezone.utc),
+        description="Timestamp of the health check",
+    )
+    version: str = Field(
+        default="1.1.0",
+        description="Server version",
+    )
+    mode: Optional[str] = Field(
+        default=None,
+        description="Instance mode: 'project' or 'shared'",
+    )
+    instance_id: Optional[str] = Field(
+        default=None,
+        description="Unique instance identifier",
+    )
+    project_id: Optional[str] = Field(
+        default=None,
+        description="Project identifier (shared mode)",
+    )
+    active_projects: Optional[int] = Field(
+        default=None,
+        description="Number of active projects (shared mode)",
+    )
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "status": "healthy",
+                    "message": "Server is running and ready for queries",
+                    "timestamp": "2024-12-15T10:30:00Z",
+                    "version": "1.1.0",
+                }
+            ]
+        }
+    }
+
+
+class IndexingStatus(BaseModel):
+    """Detailed indexing status response."""
+
+    total_documents: int = Field(
+        default=0,
+        ge=0,
+        description="Total number of documents indexed",
+    )
+    total_chunks: int = Field(
+        default=0,
+        ge=0,
+        description="Total number of chunks in vector store",
+    )
+    total_doc_chunks: int = Field(
+        default=0,
+        ge=0,
+        description="Number of document chunks",
+    )
+    total_code_chunks: int = Field(
+        default=0,
+        ge=0,
+        description="Number of code chunks",
+    )
+    supported_languages: list[str] = Field(
+        default_factory=list,
+        description="Programming languages that have been indexed",
+    )
+    indexing_in_progress: bool = Field(
+        default=False,
+        description="Whether indexing is currently in progress",
+    )
+    current_job_id: Optional[str] = Field(
+        None,
+        description="ID of the current indexing job",
+    )
+    progress_percent: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=100.0,
+        description="Progress percentage of current indexing job",
+    )
+    last_indexed_at: Optional[datetime] = Field(
+        None,
+        description="Timestamp of last completed indexing operation",
+    )
+    indexed_folders: list[str] = Field(
+        default_factory=list,
+        description="List of folders that have been indexed",
+    )
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "total_documents": 150,
+                    "total_chunks": 1200,
+                    "total_doc_chunks": 800,
+                    "total_code_chunks": 400,
+                    "indexing_in_progress": False,
+                    "current_job_id": None,
+                    "progress_percent": 0.0,
+                    "last_indexed_at": "2024-12-15T10:30:00Z",
+                    "indexed_folders": ["/path/to/docs"],
+                    "supported_languages": ["python", "typescript", "java"],
+                }
+            ]
+        }
+    }