docugraph 0.1.0__py3-none-any.whl

docugraph/agents/tools/search_code.py

@@ -0,0 +1,305 @@
+"""Search code tool for the coding agent.
+
+Provides code-specific search with language and structure awareness.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+
+def search_code(
+    query: str,
+    language: str | None = None,
+    top_k: int = 10,
+    include_context: bool = True,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search for code examples in indexed documentation, from synchronous code.
+
+    A thin wrapper over ``search_code_async`` for callers that have no event
+    loop. Async callers -- the agent loop included -- await
+    ``search_code_async`` instead: this drives a fresh loop and raises if one
+    is already running.
+
+    Args:
+        query: Search query (can be natural language or code snippet)
+        language: Filter by programming language (python, javascript, etc.)
+        top_k: Number of results to return
+        include_context: Include surrounding context
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of code results with code, path, language, and context
+    """
+    return asyncio.run(
+        search_code_async(
+            query=query,
+            language=language,
+            top_k=top_k,
+            include_context=include_context,
+            _context=_context,
+            **_kwargs,
+        )
+    )
+
+
+async def search_code_async(
+    query: str,
+    language: str | None = None,
+    top_k: int = 10,
+    include_context: bool = True,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search for code examples in indexed documentation.
+
+    This is the implementation; ``search_code`` is the synchronous wrapper.
+
+    Args:
+        query: Search query (can be natural language or code snippet)
+        language: Filter by programming language (python, javascript, etc.)
+        top_k: Number of results to return
+        include_context: Include surrounding context
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of code results with code, path, language, and context
+    """
+    from docugraph.retrieval.hybrid_search import (
+        HybridRetriever,
+        HybridSearchConfig,
+        SearchMode,
+    )
+
+    # Enhance query for code search
+    code_query = query
+    if language:
+        code_query = f"{language} {query}"
+
+    config = HybridSearchConfig(include_graph=False)
+    retriever = HybridRetriever(config=config)
+
+    results = await retriever.asearch(
+        query=code_query,
+        top_k=top_k * 2,  # Get more results for filtering
+        mode=SearchMode.HYBRID,
+    )
+
+    # Filter and format code results
+    formatted = []
+    for result in results:
+        chunk = result.chunk
+        content = chunk.content
+
+        # Skip if doesn't look like code
+        if not _contains_code(content):
+            continue
+
+        # Filter by language if specified
+        if language:
+            detected_lang = _detect_language(content)
+            if detected_lang and language.lower() not in detected_lang.lower():
+                continue
+
+        # Extract code blocks
+        code_blocks = _extract_code_blocks(content)
+        if not code_blocks:
+            # Use full content if no blocks found
+            code_blocks = [{"code": content, "language": _detect_language(content) or ""}]
+
+        for block in code_blocks:
+            formatted.append(
+                {
+                    "code": block["code"],
+                    "language": block.get("language", ""),
+                    "path": chunk.metadata.get("source_path")
+                    or chunk.metadata.get("source_url", "Unknown"),
+                    "title": chunk.metadata.get("title", ""),
+                    "score": result.score,
+                    "context": chunk.content if include_context else "",
+                    "metadata": chunk.metadata,
+                }
+            )
+
+        if len(formatted) >= top_k:
+            break
+
+    return formatted[:top_k]
+
+
+def search_symbols(
+    symbol_name: str,
+    symbol_type: str | None = None,
+    language: str | None = None,
+    top_k: int = 10,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search for code symbols (functions, classes, methods).
+
+    Args:
+        symbol_name: Name of the symbol to search for
+        symbol_type: Type filter (function, class, method, variable)
+        language: Programming language filter
+        top_k: Number of results to return
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of symbol results
+    """
+    # Build query
+    query_parts = [symbol_name]
+    if symbol_type:
+        query_parts.append(symbol_type)
+    if language:
+        query_parts.append(language)
+
+    query = " ".join(query_parts)
+
+    # Search with code focus
+    results = search_code(
+        query=query,
+        language=language,
+        top_k=top_k,
+        include_context=True,
+        _context=_context,
+    )
+
+    # Filter by symbol name presence
+    filtered = []
+    for r in results:
+        code = r.get("code", "")
+        if symbol_name.lower() in code.lower():
+            filtered.append(r)
+
+    return filtered[:top_k]
+
+
+def _contains_code(content: str) -> bool:
+    """Check if content likely contains code.
+
+    Args:
+        content: Text content to check
+
+    Returns:
+        True if content appears to contain code
+    """
+    # Check for code block markers
+    if "```" in content or "~~~" in content:
+        return True
+
+    # Check for common code patterns
+    code_indicators = [
+        "def ",
+        "class ",
+        "function ",
+        "import ",
+        "from ",
+        "const ",
+        "let ",
+        "var ",
+        "async ",
+        "await ",
+        "return ",
+        "=>",
+        "() {",
+        "(): ",
+    ]
+
+    return any(indicator in content for indicator in code_indicators)
+
+
+def _detect_language(content: str) -> str | None:
+    """Detect programming language from content.
+
+    Args:
+        content: Code content
+
+    Returns:
+        Detected language or None
+    """
+    # Check markdown code block language
+    import re
+
+    match = re.search(r"```(\w+)", content)
+    if match:
+        return match.group(1)
+
+    # Heuristic detection
+    if "def " in content and "import " in content:
+        return "python"
+    if "function " in content or "const " in content or "=>" in content:
+        return "javascript"
+    if "func " in content and "package " in content:
+        return "go"
+    if "fn " in content and "let " in content and "mut " in content:
+        return "rust"
+    if "public class " in content or "public static " in content:
+        return "java"
+
+    return None
+
+
+def _extract_code_blocks(content: str) -> list[dict[str, str]]:
+    """Extract code blocks from markdown content.
+
+    Args:
+        content: Markdown content
+
+    Returns:
+        List of code blocks with language and code
+    """
+    import re
+
+    blocks = []
+
+    # Match fenced code blocks
+    pattern = r"```(\w*)\n(.*?)```"
+    matches = re.findall(pattern, content, re.DOTALL)
+
+    for lang, code in matches:
+        blocks.append(
+            {
+                "language": lang or _detect_language(code) or "",
+                "code": code.strip(),
+            }
+        )
+
+    # Match indented code blocks (4 spaces)
+    if not blocks:
+        lines = content.split("\n")
+        code_lines = []
+        in_code = False
+
+        for line in lines:
+            if line.startswith("    ") or line.startswith("\t"):
+                code_lines.append(line[4:] if line.startswith("    ") else line[1:])
+                in_code = True
+            elif in_code and line.strip() == "":
+                code_lines.append("")
+            elif in_code:
+                if code_lines:
+                    blocks.append(
+                        {
+                            "language": _detect_language("\n".join(code_lines)) or "",
+                            "code": "\n".join(code_lines).strip(),
+                        }
+                    )
+                code_lines = []
+                in_code = False
+
+        if code_lines:
+            blocks.append(
+                {
+                    "language": _detect_language("\n".join(code_lines)) or "",
+                    "code": "\n".join(code_lines).strip(),
+                }
+            )
+
+    return blocks

docugraph/agents/tools/search_docs.py

@@ -0,0 +1,148 @@
+"""Search documentation tool for the coding agent.
+
+Provides vector and hybrid search capabilities over indexed documentation.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+
+def search_docs(
+    query: str,
+    top_k: int = 5,
+    mode: str = "hybrid",
+    include_graph: bool = False,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search indexed documentation from synchronous code.
+
+    A thin wrapper over ``search_docs_async`` for callers that have no event
+    loop. Async callers -- the agent loop included -- await
+    ``search_docs_async`` instead: this drives a fresh loop and raises if one
+    is already running.
+
+    Args:
+        query: Search query text
+        top_k: Number of results to return
+        mode: Search mode (vector, keyword, hybrid, all)
+        include_graph: Include knowledge graph in search
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of search results with content, source, and score
+    """
+    return asyncio.run(
+        search_docs_async(
+            query=query,
+            top_k=top_k,
+            mode=mode,
+            include_graph=include_graph,
+            _context=_context,
+            **_kwargs,
+        )
+    )
+
+
+async def search_docs_async(
+    query: str,
+    top_k: int = 5,
+    mode: str = "hybrid",
+    include_graph: bool = False,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search indexed documentation.
+
+    This is the implementation; ``search_docs`` is the synchronous wrapper.
+
+    Args:
+        query: Search query text
+        top_k: Number of results to return
+        mode: Search mode (vector, keyword, hybrid, all)
+        include_graph: Include knowledge graph in search
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of search results with content, source, and score
+    """
+    from docugraph.retrieval.hybrid_search import (
+        FusionStrategy,
+        HybridRetriever,
+        HybridSearchConfig,
+        SearchMode,
+    )
+
+    mode_map = {
+        "vector": SearchMode.VECTOR,
+        "keyword": SearchMode.KEYWORD,
+        "hybrid": SearchMode.HYBRID,
+        "all": SearchMode.ALL,
+    }
+
+    config = HybridSearchConfig(include_graph=include_graph)
+    retriever = HybridRetriever(config=config)
+
+    results = await retriever.asearch(
+        query=query,
+        top_k=top_k,
+        mode=mode_map.get(mode, SearchMode.HYBRID),
+        fusion=FusionStrategy.RRF,
+    )
+
+    # Format results
+    formatted = []
+    for result in results:
+        chunk = result.chunk
+        formatted.append(
+            {
+                "content": chunk.content,
+                "source": chunk.metadata.get("source_url")
+                or chunk.metadata.get("source_path", "Unknown"),
+                "title": chunk.metadata.get("title", "Untitled"),
+                "section": chunk.metadata.get("section_title", ""),
+                "score": result.score,
+                "sources": result.sources,
+                "metadata": chunk.metadata,
+            }
+        )
+
+    return formatted
+
+
+def search_by_source(
+    query: str,
+    source_pattern: str,
+    top_k: int = 10,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search documentation filtered by source pattern.
+
+    Args:
+        query: Search query text
+        source_pattern: Pattern to filter sources (e.g., "react", "fastapi")
+        top_k: Number of results to return
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of filtered search results
+    """
+    # Get more results than needed for filtering
+    results = search_docs(query=query, top_k=top_k * 3, _context=_context)
+
+    # Filter by source pattern
+    pattern_lower = source_pattern.lower()
+    filtered = [
+        r
+        for r in results
+        if pattern_lower in r.get("source", "").lower()
+        or pattern_lower in r.get("title", "").lower()
+    ]
+
+    return filtered[:top_k]

docugraph/core/__init__.py

@@ -0,0 +1,34 @@
+"""Core configuration, models, and provider abstractions."""
+
+from docugraph.core.config import Config, get_config
+from docugraph.core.embeddings import EmbeddingProviderBase, get_embedder
+from docugraph.core.llm import (
+    LLMProviderBase,
+    get_graphiti_llm_client,
+    get_llm_client,
+    get_llm_provider,
+    is_ollama_available,
+)
+from docugraph.core.models import Chunk, Document, Entity, MemoryEntry, Relationship, SearchResult
+
+__all__ = [
+    # Config
+    "Config",
+    "get_config",
+    # Models
+    "Document",
+    "Chunk",
+    "Entity",
+    "Relationship",
+    "MemoryEntry",
+    "SearchResult",
+    # Embeddings
+    "EmbeddingProviderBase",
+    "get_embedder",
+    # LLM
+    "LLMProviderBase",
+    "get_llm_provider",
+    "get_llm_client",
+    "get_graphiti_llm_client",
+    "is_ollama_available",
+]

docugraph/core/config.py

@@ -0,0 +1,183 @@
+"""Configuration management for DocuGraph AI."""
+
+import os
+from enum import StrEnum
+from functools import lru_cache
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from docugraph import _version
+
+
+class EmbeddingProvider(StrEnum):
+    """Supported embedding providers."""
+
+    AUTO = "auto"
+    FASTEMBED = "fastembed"
+    SENTENCE_TRANSFORMERS = "sentence-transformers"
+    OLLAMA = "ollama"
+    OPENAI = "openai"
+    COHERE = "cohere"
+
+
+class LLMProvider(StrEnum):
+    """Supported LLM providers."""
+
+    AUTO = "auto"
+    OLLAMA = "ollama"
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+
+
+class StorageConfig(BaseModel):
+    """Storage configuration."""
+
+    data_dir: Path = Field(default=Path.home() / ".docugraph" / "data")
+    vector_db: str = "lancedb"
+    graph_db: str = "kuzu"
+
+    def model_post_init(self, __context: Any) -> None:
+        """Ensure data directory exists."""
+        self.data_dir.mkdir(parents=True, exist_ok=True)
+
+
+class EmbeddingConfig(BaseModel):
+    """Embedding configuration (local-first, LLM-agnostic).
+
+    provider=auto resolves: fastembed (local, pure ONNX, no torch) ->
+    sentence-transformers (only if already installed, requires torch). It
+    reads nothing from the environment, so a cloud provider is used only when
+    one is named here.
+
+    ``model`` has no default, because the backends do not share one. Each local
+    backend names its own model in its constructor (fastembed's bge-small,
+    sentence-transformers' MiniLM, Ollama's nomic-embed-text); a *cloud* backend
+    has none, so leaving this unset with ``provider: openai`` or ``cohere`` is an
+    error rather than a guess. That distinction matters: a single shared default
+    here was the local fastembed name, so naming a cloud provider without naming
+    a model would have sent ``BAAI/bge-small-en-v1.5`` to an API that has never
+    heard of it.
+    """
+
+    provider: EmbeddingProvider = EmbeddingProvider.AUTO
+    model: str | None = None  # local backends have their own; cloud has none
+    dimensions: int | None = None  # measured from the model unless set here
+    batch_size: int = 32
+
+
+class LLMConfig(BaseModel):
+    """LLM configuration for extraction and reranking.
+
+    ``model`` defaults to the ``"auto"`` sentinel, which means "whichever model
+    this provider offers". Only a local Ollama server can answer that, because
+    it is the only provider docugraph can ask what it has installed. A cloud
+    provider needs the model named: every name that could be written here has a
+    retirement date on the provider's calendar, and a retired name fails every
+    call with a 404 the user cannot act on.
+
+    ``temperature`` is unset by default and is sent only when set. That is not
+    tidiness: OpenAI's reasoning models reject any temperature other than their
+    fixed default, so a default of 0.0 here made every call to one fail with
+    ``unsupported value``. Set ``temperature: 0.0`` explicitly for deterministic
+    local output.
+    """
+
+    provider: LLMProvider = LLMProvider.AUTO
+    model: str = "auto"
+    # Graphiti asks for two models: a large one for extraction and a small one
+    # for the cheaper summarization and edge-dedup passes. Unset means "the same
+    # as `model`, whatever that turns out to be" -- not a hard-coded model name,
+    # which is what it was.
+    small_model: str | None = None
+    temperature: float | None = None
+    max_tokens: int = 4096
+
+    # API keys (loaded from environment)
+    openai_api_key: str | None = None
+    anthropic_api_key: str | None = None
+
+    def model_post_init(self, __context: Any) -> None:
+        """Load API keys from environment."""
+        self.openai_api_key = os.getenv("OPENAI_API_KEY")
+        self.anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
+
+
+class CrawlerConfig(BaseModel):
+    """Web crawler configuration."""
+
+    max_concurrent: int = 5
+    rate_limit: float = 2.0  # requests per second
+    user_agent: str = f"DocuGraph-AI/{_version.__version__}"
+    respect_robots: bool = True
+    cache_ttl: int = 86400  # 24 hours
+    timeout: int = 30  # seconds
+
+
+class Config(BaseSettings):
+    """Main configuration for DocuGraph AI."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="DOCUGRAPH_",
+        env_nested_delimiter="__",
+        extra="ignore",
+    )
+
+    storage: StorageConfig = Field(default_factory=StorageConfig)
+    embeddings: EmbeddingConfig = Field(default_factory=EmbeddingConfig)
+    llm: LLMConfig = Field(default_factory=LLMConfig)
+    crawler: CrawlerConfig = Field(default_factory=CrawlerConfig)
+
+    @classmethod
+    def from_yaml(cls, path: Path) -> "Config":
+        """Load configuration from a YAML file."""
+        if not path.exists():
+            return cls()
+
+        with open(path) as f:
+            data = yaml.safe_load(f) or {}
+
+        return cls(**data)
+
+    @classmethod
+    def load(cls) -> "Config":
+        """Load configuration from default locations."""
+        # Check for config file in order of precedence
+        config_paths = [
+            Path.cwd() / "docugraph.yaml",
+            Path.cwd() / "docugraph.yml",
+            Path.cwd() / ".docugraph.yaml",
+            Path.home() / ".docugraph" / "config.yaml",
+            Path.home() / ".config" / "docugraph" / "config.yaml",
+        ]
+
+        for path in config_paths:
+            if path.exists():
+                return cls.from_yaml(path)
+
+        return cls()
+
+    def save(self, path: Path | None = None) -> None:
+        """Save configuration to a YAML file."""
+        if path is None:
+            path = Path.home() / ".docugraph" / "config.yaml"
+
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        data = self.model_dump(mode="json")
+        with open(path, "w") as f:
+            yaml.dump(data, f, default_flow_style=False, sort_keys=False)
+
+
+@lru_cache
+def get_config() -> Config:
+    """Get the singleton configuration instance."""
+    return Config.load()
+
+
+def reset_config() -> None:
+    """Reset the cached configuration (useful for testing)."""
+    get_config.cache_clear()