docugraph 0.1.0__py3-none-any.whl

docugraph/agents/tools/graph_query.py

@@ -0,0 +1,203 @@
+"""Knowledge graph query tool for the coding agent.
+
+Provides access to the knowledge graph for entity and relationship queries.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+
+async def graph_query(
+    query: str,
+    num_results: int = 10,
+    group_id: str | None = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Query the knowledge graph for facts and relationships.
+
+    Args:
+        query: Natural language query
+        num_results: Maximum number of results
+        group_id: Optional group ID filter
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of facts from the knowledge graph
+    """
+    from docugraph.storage.graph_store import GraphStore
+
+    store = GraphStore()
+    group_ids = [group_id] if group_id else None
+
+    try:
+        results = await store.search(
+            query=query,
+            group_ids=group_ids,
+            num_results=num_results,
+        )
+        return results
+    finally:
+        await store.close()
+
+
+async def graph_add_content(
+    content: str,
+    name: str | None = None,
+    group_id: str = "default",
+    source_type: str = "text",
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any]:
+    """Add content to the knowledge graph.
+
+    Args:
+        content: Content to extract entities from
+        name: Episode name (optional)
+        group_id: Group ID for scoping
+        source_type: Content type (text, json)
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Information about extracted entities and relationships
+    """
+    from docugraph.storage.graph_store import GraphStore
+
+    store = GraphStore()
+
+    try:
+        episode_name = name or f"episode_{hash(content) % 10000}"
+        result = await store.add_episode(
+            name=episode_name,
+            content=content,
+            source_type=source_type,
+            group_id=group_id,
+        )
+        return result
+    finally:
+        await store.close()
+
+
+async def find_related_entities(
+    entity_name: str,
+    relationship_type: str | None = None,
+    max_depth: int = 2,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Find entities related to a given entity.
+
+    Args:
+        entity_name: Name of the entity to start from
+        relationship_type: Filter by relationship type (USES, DEPENDS_ON, etc.)
+        max_depth: How deep to look — widens the result pool per hop
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of related entities with relationship information
+    """
+    # Query for the entity and its relationships
+    query = f"What is related to {entity_name}?"
+    if relationship_type:
+        query = f"What does {entity_name} {relationship_type.lower().replace('_', ' ')}?"
+
+    # Deeper searches pull a proportionally wider pool of graph facts.
+    results = await graph_query(
+        query=query,
+        num_results=10 * max(1, max_depth),
+        _context=_context,
+    )
+
+    # Filter results that mention the entity
+    related = []
+    entity_lower = entity_name.lower()
+
+    for r in results:
+        fact = r.get("fact", "").lower()
+        if entity_lower in fact:
+            related.append(r)
+
+    return related
+
+
+async def get_entity_context(
+    entity_name: str,
+    include_relationships: bool = True,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any]:
+    """Get comprehensive context about an entity.
+
+    Args:
+        entity_name: Name of the entity
+        include_relationships: Include relationship information
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Entity context including facts, relationships, and documentation
+    """
+    # Query for entity facts
+    entity_facts = await graph_query(
+        query=f"What is {entity_name}?",
+        num_results=10,
+        _context=_context,
+    )
+
+    context = {
+        "entity": entity_name,
+        "facts": entity_facts,
+        "relationships": [],
+        "documentation": [],
+    }
+
+    if include_relationships:
+        # Get relationships
+        relationships = await find_related_entities(
+            entity_name=entity_name,
+            _context=_context,
+        )
+        context["relationships"] = relationships
+
+    # Also search documentation. This function is async, so it awaits the
+    # async entry point rather than driving a loop of its own.
+    from docugraph.agents.tools.search_docs import search_docs_async
+
+    doc_results = await search_docs_async(
+        query=entity_name,
+        top_k=3,
+        _context=_context,
+    )
+    context["documentation"] = doc_results
+
+    return context
+
+
+def query_entities_sync(
+    query: str,
+    num_results: int = 10,
+    group_id: str | None = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Synchronous wrapper for graph_query, for callers with no event loop.
+
+    Async callers await ``graph_query`` itself -- this drives a fresh loop and
+    raises if one is already running, which is the signal to await instead.
+
+    Args:
+        query: Natural language query
+        num_results: Maximum number of results
+        group_id: Optional group ID filter
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of facts from the knowledge graph
+    """
+    return asyncio.run(graph_query(query, num_results, group_id, _context))

docugraph/agents/tools/memory_ops.py

@@ -0,0 +1,305 @@
+"""Memory operations tool for the coding agent.
+
+Provides session memory for storing and retrieving context, decisions,
+and other information scoped by repository and branch.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import Any
+
+
+def _scope_or_defaults(
+    repository: str | None,
+    branch: str | None,
+    _context: Any,
+) -> tuple[str, str]:
+    """Resolve repository/branch scope: explicit arg, then context, then defaults."""
+    repo = repository or (getattr(_context, "repository", "default") if _context else "default")
+    br = branch or (getattr(_context, "branch", "main") if _context else "main")
+    return str(repo), str(br)
+
+
+def memory_store(
+    key: str,
+    value: Any,
+    repository: str | None = None,
+    branch: str | None = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any]:
+    """Store a value in agent memory.
+
+    Args:
+        key: Key to store the value under
+        value: Value to store (can be any JSON-serializable type)
+        repository: Repository scope (defaults to context or 'default')
+        branch: Branch scope (defaults to context or 'main')
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Confirmation with stored entry details
+    """
+    from docugraph.storage.memory_store import MemoryStore
+
+    # Get scope from context or defaults
+    repository, branch = _scope_or_defaults(repository, branch, _context)
+
+    store = MemoryStore()
+    entry = store.set(
+        key=key,
+        value=value,
+        repository=repository,
+        branch=branch,
+    )
+
+    return {
+        "key": entry.key,
+        "value": entry.value,
+        "repository": entry.repository,
+        "branch": entry.branch,
+        "created_at": entry.created_at.isoformat() if entry.created_at else None,
+    }
+
+
+def memory_recall(
+    key: str,
+    repository: str | None = None,
+    branch: str | None = None,
+    default: Any = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> Any:
+    """Retrieve a value from agent memory.
+
+    Args:
+        key: Key to retrieve (or '*' for all keys)
+        repository: Repository scope (defaults to context or 'default')
+        branch: Branch scope (defaults to context or 'main')
+        default: Default value if key not found
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Stored value or default if not found
+    """
+    from docugraph.storage.memory_store import MemoryStore
+
+    # Get scope from context or defaults
+    repository, branch = _scope_or_defaults(repository, branch, _context)
+
+    store = MemoryStore()
+
+    # List all keys
+    if key == "*":
+        keys = store.list_keys(repository=repository, branch=branch)
+        return {
+            "keys": keys,
+            "repository": repository,
+            "branch": branch,
+            "total": len(keys),
+        }
+
+    # Get specific key
+    value = store.get(key, repository=repository, branch=branch)
+    if value is None:
+        return default
+
+    return value
+
+
+def memory_delete(
+    key: str,
+    repository: str | None = None,
+    branch: str | None = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any]:
+    """Delete a value from agent memory.
+
+    Args:
+        key: Key to delete
+        repository: Repository scope
+        branch: Branch scope
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Confirmation with deletion status
+    """
+    from docugraph.storage.memory_store import MemoryStore
+
+    # Get scope from context or defaults
+    repository, branch = _scope_or_defaults(repository, branch, _context)
+
+    store = MemoryStore()
+    deleted = store.delete(key, repository=repository, branch=branch)
+
+    return {
+        "key": key,
+        "deleted": deleted,
+        "repository": repository,
+        "branch": branch,
+    }
+
+
+def memory_search(
+    pattern: str,
+    repository: str | None = None,
+    branch: str | None = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Search memory keys by pattern.
+
+    Args:
+        pattern: Pattern to search for (supports % wildcard)
+        repository: Repository scope
+        branch: Branch scope
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of matching entries with keys and values
+    """
+    from docugraph.storage.memory_store import MemoryStore
+
+    # Get scope from context or defaults
+    repository, branch = _scope_or_defaults(repository, branch, _context)
+
+    store = MemoryStore()
+    keys = store.list_keys(repository=repository, branch=branch)
+
+    # Filter by pattern
+    import fnmatch
+
+    # Convert SQL-style wildcard to glob-style
+    glob_pattern = pattern.replace("%", "*")
+    matching_keys = fnmatch.filter(keys, glob_pattern)
+
+    # Get values for matching keys
+    results = []
+    for key in matching_keys:
+        value = store.get(key, repository=repository, branch=branch)
+        results.append(
+            {
+                "key": key,
+                "value": value,
+                "repository": repository,
+                "branch": branch,
+            }
+        )
+
+    return results
+
+
+def store_decision(
+    decision: str,
+    reasoning: str,
+    context: dict[str, Any] | None = None,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any]:
+    """Store a decision with reasoning for future reference.
+
+    Args:
+        decision: The decision made
+        reasoning: Why this decision was made
+        context: Additional context
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Stored decision entry
+    """
+    # Generate a unique key for the decision
+    timestamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S")
+    key = f"decision_{timestamp}"
+
+    value = {
+        "decision": decision,
+        "reasoning": reasoning,
+        "context": context or {},
+        "timestamp": datetime.now(UTC).isoformat(),
+    }
+
+    return memory_store(key=key, value=value, _context=_context)
+
+
+def recall_decisions(
+    limit: int = 10,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> list[dict[str, Any]]:
+    """Recall recent decisions.
+
+    Args:
+        limit: Maximum number of decisions to return
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        List of recent decisions
+    """
+    results = memory_search(pattern="decision_%", _context=_context)
+
+    # Sort by key (which includes timestamp) and return most recent
+    results.sort(key=lambda x: x["key"], reverse=True)
+
+    decisions = []
+    for r in results[:limit]:
+        value = r.get("value", {})
+        if isinstance(value, dict):
+            decisions.append(
+                {
+                    "key": r["key"],
+                    "decision": value.get("decision"),
+                    "reasoning": value.get("reasoning"),
+                    "timestamp": value.get("timestamp"),
+                }
+            )
+
+    return decisions
+
+
+def store_context(
+    context_type: str,
+    data: dict[str, Any],
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any]:
+    """Store contextual information.
+
+    Args:
+        context_type: Type of context (e.g., 'architecture', 'dependencies')
+        data: Context data
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Stored context entry
+    """
+    key = f"context_{context_type}"
+    return memory_store(key=key, value=data, _context=_context)
+
+
+def recall_context(
+    context_type: str,
+    _context: Any = None,
+    **_kwargs: Any,
+) -> dict[str, Any] | None:
+    """Recall contextual information.
+
+    Args:
+        context_type: Type of context to recall
+        _context: Agent execution context (injected)
+        **kwargs: Additional arguments
+
+    Returns:
+        Context data or None if not found
+    """
+    key = f"context_{context_type}"
+    result: dict[str, Any] | None = memory_recall(key=key, default=None, _context=_context)
+    return result