kg-mcp 0.1.8__py3-none-any.whl

kg_mcp/codegraph/model.py

@@ -0,0 +1,170 @@
+"""
+Data models for code graph entities.
+Represents files, symbols, and their relationships.
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import List, Optional
+
+
+class SymbolKind(str, Enum):
+    """Types of code symbols."""
+
+    FILE = "file"
+    MODULE = "module"
+    CLASS = "class"
+    FUNCTION = "function"
+    METHOD = "method"
+    PROPERTY = "property"
+    VARIABLE = "variable"
+    CONSTANT = "constant"
+    INTERFACE = "interface"
+    ENUM = "enum"
+    TYPE_ALIAS = "type_alias"
+
+
+class ReferenceKind(str, Enum):
+    """Types of symbol references."""
+
+    CALL = "call"
+    IMPORT = "import"
+    INHERIT = "inherit"
+    IMPLEMENT = "implement"
+    USE = "use"
+    OVERRIDE = "override"
+
+
+@dataclass
+class SourceLocation:
+    """Location in source code."""
+
+    file_path: str
+    start_line: int
+    start_column: int = 0
+    end_line: Optional[int] = None
+    end_column: Optional[int] = None
+
+    def __post_init__(self):
+        if self.end_line is None:
+            self.end_line = self.start_line
+
+
+@dataclass
+class Symbol:
+    """Represents a code symbol (function, class, variable, etc.)."""
+
+    fqn: str  # Fully qualified name
+    name: str  # Short name
+    kind: SymbolKind
+    location: SourceLocation
+    signature: Optional[str] = None
+    docstring: Optional[str] = None
+    parent_fqn: Optional[str] = None  # Parent symbol (e.g., class for method)
+    modifiers: List[str] = field(default_factory=list)  # public, private, static, async, etc.
+
+    @property
+    def file_path(self) -> str:
+        return self.location.file_path
+
+
+@dataclass
+class SymbolReference:
+    """Represents a reference from one symbol to another."""
+
+    source_fqn: str  # Symbol making the reference
+    target_fqn: str  # Symbol being referenced
+    kind: ReferenceKind
+    location: SourceLocation
+    context: Optional[str] = None  # Line of code with reference
+
+
+@dataclass
+class FileInfo:
+    """Metadata about a source file."""
+
+    path: str
+    language: str
+    content_hash: str
+    size_bytes: int
+    line_count: int
+    last_modified: datetime
+    git_commit: Optional[str] = None
+    symbols: List[Symbol] = field(default_factory=list)
+
+    def add_symbol(self, symbol: Symbol) -> None:
+        """Add a symbol to this file."""
+        self.symbols.append(symbol)
+
+
+@dataclass
+class CodeGraphSnapshot:
+    """A snapshot of the entire code graph."""
+
+    project_id: str
+    timestamp: datetime
+    files: List[FileInfo]
+    references: List[SymbolReference]
+
+    @property
+    def total_symbols(self) -> int:
+        return sum(len(f.symbols) for f in self.files)
+
+    @property
+    def total_files(self) -> int:
+        return len(self.files)
+
+    @property
+    def total_references(self) -> int:
+        return len(self.references)
+
+
+# Language detection based on file extension
+LANGUAGE_EXTENSIONS = {
+    ".py": "python",
+    ".pyi": "python",
+    ".js": "javascript",
+    ".jsx": "javascript",
+    ".ts": "typescript",
+    ".tsx": "typescript",
+    ".java": "java",
+    ".kt": "kotlin",
+    ".go": "go",
+    ".rs": "rust",
+    ".c": "c",
+    ".cpp": "cpp",
+    ".cc": "cpp",
+    ".h": "c",
+    ".hpp": "cpp",
+    ".cs": "csharp",
+    ".rb": "ruby",
+    ".php": "php",
+    ".swift": "swift",
+    ".scala": "scala",
+    ".r": "r",
+    ".R": "r",
+    ".sql": "sql",
+    ".sh": "bash",
+    ".bash": "bash",
+    ".zsh": "zsh",
+    ".yaml": "yaml",
+    ".yml": "yaml",
+    ".json": "json",
+    ".toml": "toml",
+    ".xml": "xml",
+    ".html": "html",
+    ".css": "css",
+    ".scss": "scss",
+    ".less": "less",
+    ".md": "markdown",
+    ".rst": "rst",
+}
+
+
+def detect_language(file_path: str) -> str:
+    """Detect language from file extension."""
+    from pathlib import Path
+
+    ext = Path(file_path).suffix.lower()
+    return LANGUAGE_EXTENSIONS.get(ext, "unknown")

kg_mcp/config.py

@@ -0,0 +1,83 @@
+"""
+Configuration management for MCP-KG-Memory server.
+Uses pydantic-settings for environment variable handling.
+"""
+
+from functools import lru_cache
+from typing import List
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    """Application settings loaded from environment variables."""
+
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        case_sensitive=False,
+        extra="ignore",
+    )
+
+    # Neo4j Configuration
+    neo4j_uri: str = Field(default="bolt://localhost:7687", description="Neo4j Bolt URI")
+    neo4j_user: str = Field(default="neo4j", description="Neo4j username")
+    neo4j_password: str = Field(default="password123", description="Neo4j password")
+    neo4j_configured: str = Field(default="1", description="Is Neo4j configured (1/0)")
+
+    # LLM Configuration (supports both direct Gemini and LiteLLM Gateway)
+    # Mode: 'gemini_direct', 'litellm', or 'both'
+    llm_mode: str = Field(default="litellm", description="Operation mode")
+    llm_primary: str = Field(default="litellm", description="Primary provider if both configured")
+    llm_provider: str = Field(default="litellm", description="Active provider tag")
+
+    # Gemini Direct
+    gemini_api_key: str = Field(default="", description="Gemini API Key (direct)")
+    gemini_base_url: str = Field(default="https://generativelanguage.googleapis.com/", description="Gemini API Base URL")
+    gemini_model: str = Field(default="", description="Gemini Direct Model ID")
+
+    # LiteLLM
+    litellm_api_key: str = Field(default="", description="LiteLLM Gateway API Key")
+    litellm_base_url: str = Field(default="", description="LiteLLM Gateway Base URL")
+    litellm_model: str = Field(default="", description="LiteLLM Model ID")
+
+    # General / Fallback Model
+    llm_model: str = Field(
+        default="gemini/gemini-1.5-flash",
+        description="Default model identifier (legacy fallback)",
+    )
+    
+    # Task-specific Routing (optional)
+    kg_model_default: str = Field(default="", description="Default model for general tasks")
+    kg_model_fast: str = Field(default="", description="Fast model for high-throughput")
+    kg_model_reason: str = Field(default="", description="Reasoning model for complex tasks")
+
+    llm_temperature: float = Field(default=0.2, description="LLM temperature for extraction")
+    llm_max_tokens: int = Field(default=4096, description="Maximum tokens for LLM response")
+
+    # MCP Server Configuration
+    mcp_host: str = Field(default="127.0.0.1", description="MCP server host")
+    mcp_port: int = Field(default=8000, description="MCP server port")
+    mcp_stateless: bool = Field(default=True, description="Run server in stateless mode")
+
+    # Security Configuration
+    kg_mcp_token: str = Field(default="", description="Bearer token for authentication")
+    kg_allowed_origins: str = Field(
+        default="http://localhost:*,http://127.0.0.1:*",
+        description="Comma-separated list of allowed origins",
+    )
+
+    # Logging
+    log_level: str = Field(default="INFO", description="Logging level")
+
+    @property
+    def allowed_origins_list(self) -> List[str]:
+        """Parse allowed origins into a list."""
+        return [origin.strip() for origin in self.kg_allowed_origins.split(",") if origin.strip()]
+
+
+@lru_cache
+def get_settings() -> Settings:
+    """Get cached settings instance."""
+    return Settings()

kg_mcp/kg/__init__.py

@@ -0,0 +1,3 @@
+"""
+Knowledge Graph submodule for Neo4j operations.
+"""

kg_mcp/kg/apply_schema.py

@@ -0,0 +1,93 @@
+"""
+Apply Neo4j schema constraints and indexes.
+Run this script to initialize the database schema.
+"""
+
+import asyncio
+import logging
+import sys
+from pathlib import Path
+
+from kg_mcp.kg.neo4j import get_neo4j_client, init_neo4j, close_neo4j
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+async def apply_schema() -> None:
+    """Read and apply schema.cypher to Neo4j."""
+    logger.info("Connecting to Neo4j...")
+    await init_neo4j()
+
+    client = get_neo4j_client()
+
+    # Read schema file
+    schema_path = Path(__file__).parent / "schema.cypher"
+    if not schema_path.exists():
+        logger.error(f"Schema file not found: {schema_path}")
+        sys.exit(1)
+
+    logger.info(f"Reading schema from {schema_path}")
+    schema_content = schema_path.read_text()
+
+    # Split into individual statements (skip comments and empty lines)
+    statements = []
+    for line in schema_content.split("\n"):
+        line = line.strip()
+        # Skip comments and empty lines
+        if not line or line.startswith("//"):
+            continue
+        statements.append(line)
+
+    # Join multi-line statements
+    full_statements = []
+    current_stmt = []
+    for line in statements:
+        current_stmt.append(line)
+        if line.endswith(";"):
+            full_stmt = " ".join(current_stmt).rstrip(";")
+            full_statements.append(full_stmt)
+            current_stmt = []
+
+    logger.info(f"Found {len(full_statements)} schema statements")
+
+    # Execute each statement
+    success_count = 0
+    error_count = 0
+
+    for i, stmt in enumerate(full_statements, 1):
+        if not stmt.strip():
+            continue
+
+        try:
+            logger.debug(f"Executing statement {i}: {stmt[:80]}...")
+            await client.execute_query(stmt)
+            success_count += 1
+            logger.info(f"✓ Statement {i} applied successfully")
+        except Exception as e:
+            error_msg = str(e)
+            # Ignore "already exists" errors for constraints/indexes
+            if "already exists" in error_msg.lower() or "equivalent" in error_msg.lower():
+                logger.info(f"⊘ Statement {i} skipped (already exists)")
+                success_count += 1
+            else:
+                logger.error(f"✗ Statement {i} failed: {e}")
+                error_count += 1
+
+    logger.info(f"\nSchema application complete:")
+    logger.info(f"  ✓ Success: {success_count}")
+    logger.info(f"  ✗ Errors:  {error_count}")
+
+    await close_neo4j()
+
+
+def main() -> None:
+    """Entry point for schema application."""
+    asyncio.run(apply_schema())
+
+
+if __name__ == "__main__":
+    main()

kg_mcp/kg/ingest.py

@@ -0,0 +1,253 @@
+"""
+Ingest pipeline for processing user requests.
+Orchestrates LLM extraction, linking, and Neo4j commit.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from kg_mcp.llm.client import get_llm_client
+from kg_mcp.llm.schemas import ExtractionResult, LinkingResult
+from kg_mcp.kg.repo import get_repository
+
+logger = logging.getLogger(__name__)
+
+
+class IngestPipeline:
+    """Pipeline for ingesting user interactions into the knowledge graph."""
+
+    def __init__(self):
+        self.llm = get_llm_client()
+        self.repo = get_repository()
+
+    async def process_message(
+        self,
+        project_id: str,
+        user_text: str,
+        user_id: str = "default_user",
+        files: Optional[List[str]] = None,
+        diff: Optional[str] = None,
+        symbols: Optional[List[str]] = None,
+        tags: Optional[List[str]] = None,
+    ) -> Dict[str, Any]:
+        """
+        Process a user message through the full pipeline.
+
+        Args:
+            project_id: Project ID to associate with this interaction
+            user_text: The user's message text
+            user_id: User ID for preferences
+            files: Optional list of file paths involved
+            diff: Optional code diff
+            symbols: Optional list of code symbols
+            tags: Optional tags for this interaction
+
+        Returns:
+            Dict containing interaction_id, extracted entities, and created entity IDs
+        """
+        logger.info(f"Processing message for project {project_id}")
+
+        # Step 0: Ensure project exists
+        await self.repo.get_or_create_project(project_id)
+
+        # Step 1: Create interaction record
+        interaction = await self.repo.create_interaction(
+            project_id=project_id,
+            user_text=user_text,
+            tags=tags,
+        )
+        interaction_id = interaction["id"]
+        logger.info(f"Created interaction {interaction_id}")
+
+        # Step 2: Extract entities using LLM
+        extraction = await self.llm.extract_entities(
+            user_text=user_text,
+            files=files,
+            diff=diff,
+            symbols=symbols,
+        )
+        logger.info(
+            f"Extracted: {len(extraction.goals)} goals, "
+            f"{len(extraction.constraints)} constraints, "
+            f"{len(extraction.preferences)} preferences, "
+            f"{len(extraction.pain_points)} pain points, "
+            f"{len(extraction.strategies)} strategies"
+        )
+
+        # Step 3: Get existing entities for linking
+        existing_goals = await self.repo.get_all_goals(project_id)
+        existing_preferences = await self.repo.get_preferences(user_id)
+        recent_interactions = await self.repo.get_recent_interactions(project_id, limit=5)
+
+        # Step 4: Link entities using LLM
+        linking = await self.llm.link_entities(
+            extraction=extraction,
+            existing_goals=existing_goals,
+            existing_preferences=existing_preferences,
+            recent_interactions=recent_interactions,
+        )
+        logger.info(
+            f"Linking: {len(linking.merge_suggestions)} merges, "
+            f"{len(linking.relationships)} relationships"
+        )
+
+        # Step 5: Commit to Neo4j
+        created_entities = await self._commit_to_graph(
+            project_id=project_id,
+            user_id=user_id,
+            interaction_id=interaction_id,
+            extraction=extraction,
+            linking=linking,
+        )
+
+        return {
+            "interaction_id": interaction_id,
+            "extracted": extraction.model_dump(),
+            "linking": linking.model_dump(),
+            "created_entities": created_entities,
+            "confidence": extraction.confidence,
+        }
+
+    async def _commit_to_graph(
+        self,
+        project_id: str,
+        user_id: str,
+        interaction_id: str,
+        extraction: ExtractionResult,
+        linking: LinkingResult,
+    ) -> Dict[str, List[str]]:
+        """
+        Commit extracted entities to Neo4j.
+
+        Returns dict of entity type -> list of created IDs.
+        """
+        created = {
+            "goals": [],
+            "constraints": [],
+            "preferences": [],
+            "pain_points": [],
+            "strategies": [],
+            "code_artifacts": [],
+        }
+
+        # Process merge suggestions first to build ID mapping
+        merge_map: Dict[str, str] = {}  # new_title -> existing_id
+        for merge in linking.merge_suggestions:
+            if merge.confidence >= 0.7:
+                merge_map[merge.new_entity_title] = merge.existing_entity_id
+                logger.info(
+                    f"Merging '{merge.new_entity_title}' into existing "
+                    f"'{merge.existing_entity_title}'"
+                )
+
+        # Create/update goals
+        goal_id_map: Dict[str, str] = {}  # title -> id
+        for goal_extract in extraction.goals:
+            if goal_extract.title in merge_map:
+                goal_id = merge_map[goal_extract.title]
+                goal_id_map[goal_extract.title] = goal_id
+            else:
+                goal = await self.repo.upsert_goal(
+                    project_id=project_id,
+                    title=goal_extract.title,
+                    description=goal_extract.description,
+                    status=goal_extract.status,
+                    priority=goal_extract.priority,
+                )
+                goal_id = goal["id"]
+                goal_id_map[goal_extract.title] = goal_id
+                created["goals"].append(goal_id)
+
+            # Link interaction to goal
+            await self.repo.link_interaction_to_goal(interaction_id, goal_id)
+
+        # Create/update constraints
+        for constraint_extract in extraction.constraints:
+            # Find related goal if mentioned
+            related_goal_id = None
+            for goal_extract in extraction.goals:
+                if goal_extract.title in goal_id_map:
+                    related_goal_id = goal_id_map[goal_extract.title]
+                    break
+
+            constraint = await self.repo.upsert_constraint(
+                project_id=project_id,
+                constraint_type=constraint_extract.type,
+                description=constraint_extract.description,
+                severity=constraint_extract.severity,
+                goal_id=related_goal_id,
+            )
+            created["constraints"].append(constraint["id"])
+
+        # Create/update preferences
+        for pref_extract in extraction.preferences:
+            pref = await self.repo.upsert_preference(
+                user_id=user_id,
+                category=pref_extract.category,
+                preference=pref_extract.preference,
+                strength=pref_extract.strength,
+            )
+            created["preferences"].append(pref["id"])
+
+        # Create/update pain points
+        for pp_extract in extraction.pain_points:
+            related_goal_id = None
+            if pp_extract.related_goal and pp_extract.related_goal in goal_id_map:
+                related_goal_id = goal_id_map[pp_extract.related_goal]
+
+            pp = await self.repo.upsert_painpoint(
+                project_id=project_id,
+                description=pp_extract.description,
+                severity=pp_extract.severity,
+                related_goal_id=related_goal_id,
+                interaction_id=interaction_id,
+            )
+            created["pain_points"].append(pp["id"])
+
+        # Create/update strategies
+        for strategy_extract in extraction.strategies:
+            related_goal_id = None
+            if strategy_extract.related_goal and strategy_extract.related_goal in goal_id_map:
+                related_goal_id = goal_id_map[strategy_extract.related_goal]
+
+            strategy = await self.repo.upsert_strategy(
+                project_id=project_id,
+                title=strategy_extract.title,
+                approach=strategy_extract.approach,
+                rationale=strategy_extract.rationale,
+                outcome=strategy_extract.outcome,
+                outcome_reason=strategy_extract.outcome_reason,
+                related_goal_id=related_goal_id,
+            )
+            created["strategies"].append(strategy["id"])
+
+        # Create code references as artifacts
+        for code_ref in extraction.code_references:
+            # Find related goals
+            related_goal_ids = list(goal_id_map.values())[:3]  # Link to first 3 goals
+
+            artifact = await self.repo.upsert_code_artifact(
+                project_id=project_id,
+                path=code_ref.path,
+                kind="file",
+                symbol_fqn=code_ref.symbol,
+                start_line=code_ref.start_line,
+                end_line=code_ref.end_line,
+                related_goal_ids=related_goal_ids if related_goal_ids else None,
+            )
+            created["code_artifacts"].append(artifact["id"])
+
+        logger.info(f"Committed to graph: {created}")
+        return created
+
+
+# Factory function
+_pipeline: Optional[IngestPipeline] = None
+
+
+def get_ingest_pipeline() -> IngestPipeline:
+    """Get or create the ingest pipeline singleton."""
+    global _pipeline
+    if _pipeline is None:
+        _pipeline = IngestPipeline()
+    return _pipeline