kg-mcp 0.1.8__py3-none-any.whl

kg_mcp/llm/prompts/__init__.py

@@ -0,0 +1,8 @@
+"""
+Prompt templates for LLM operations.
+"""
+
+from kg_mcp.llm.prompts.extractor import get_extractor_prompt
+from kg_mcp.llm.prompts.linker import get_linker_prompt
+
+__all__ = ["get_extractor_prompt", "get_linker_prompt"]

kg_mcp/llm/prompts/extractor.py

@@ -0,0 +1,84 @@
+"""
+Extractor prompt template for entity extraction from user text.
+"""
+
+from typing import List, Optional, Tuple
+
+
+EXTRACTOR_SYSTEM_PROMPT = """You are an expert at analyzing developer requests and extracting structured information.
+Your task is to analyze the user's message and extract:
+1. GOALS: What the user wants to achieve (objectives, features, fixes)
+2. CONSTRAINTS: Limitations or requirements (budget, technology stack, performance, time)
+3. PREFERENCES: User's coding/architecture preferences (patterns, styles, tools)
+4. PAIN_POINTS: Problems or frustrations mentioned
+5. STRATEGIES: Approaches or plans mentioned for solving problems (track Success/Failure outcomes!)
+6. ACCEPTANCE_CRITERIA: Success conditions mentioned
+7. CODE_REFERENCES: Any files, functions, classes, or code snippets referenced
+8. NEXT_ACTIONS: Immediate next steps implied by the request
+
+IMPORTANT RULES:
+- Only extract information that is explicitly stated or strongly implied
+- Be precise and concise in descriptions
+- For code_references, extract paths exactly as mentioned
+- Set confidence based on how clear the extraction is (0.0 to 1.0)
+- If the message is just a greeting or acknowledgment, return empty arrays
+- **CRITICAL**: If a strategy is mentioned as successful or failed, mark the `outcome` field!
+
+OUTPUT FORMAT (JSON):
+{
+  "goals": [{"title": "...", "description": "...", "priority": 1-5, "status": "active", "parent_goal_title": null}],
+  "constraints": [{"type": "budget|stack|style|performance|time", "description": "...", "severity": "must|should|nice_to_have"}],
+  "preferences": [{"category": "coding_style|architecture|testing|tools|output_format", "preference": "...", "strength": "prefer|avoid|require"}],
+  "pain_points": [{"description": "...", "severity": "low|medium|high|critical", "related_goal": null}],
+  "strategies": [{"title": "...", "approach": "...", "rationale": "...", "outcome": "success|failure|pending", "outcome_reason": "...", "related_goal": null}],
+  "acceptance_criteria": [{"criterion": "...", "related_goal": "...", "testable": true}],
+  "code_references": [{"path": "...", "symbol": null, "start_line": null, "end_line": null, "action": "reference|create|modify|delete"}],
+  "next_actions": ["action 1", "action 2"],
+  "confidence": 0.85
+}"""
+
+
+def get_extractor_prompt(
+    user_text: str,
+    files: Optional[List[str]] = None,
+    diff: Optional[str] = None,
+    symbols: Optional[List[str]] = None,
+    context: Optional[str] = None,
+) -> Tuple[str, str]:
+    """
+    Build the extractor prompt for entity extraction.
+
+    Args:
+        user_text: The user's message
+        files: Optional list of files involved
+        diff: Optional code diff
+        symbols: Optional list of symbols
+        context: Optional additional context
+
+    Returns:
+        Tuple of (system_prompt, user_prompt)
+    """
+    user_prompt_parts = [f"USER MESSAGE:\n{user_text}"]
+
+    if files:
+        files_str = "\n".join(f"- {f}" for f in files)
+        user_prompt_parts.append(f"\nFILES INVOLVED:\n{files_str}")
+
+    if diff:
+        # Truncate large diffs
+        truncated_diff = diff[:2000] + "..." if len(diff) > 2000 else diff
+        user_prompt_parts.append(f"\nCODE DIFF:\n```\n{truncated_diff}\n```")
+
+    if symbols:
+        symbols_str = ", ".join(symbols)
+        user_prompt_parts.append(f"\nSYMBOLS REFERENCED: {symbols_str}")
+
+    if context:
+        user_prompt_parts.append(f"\nADDITIONAL CONTEXT:\n{context}")
+
+    user_prompt_parts.append(
+        "\n\nAnalyze the above and extract structured information. "
+        "Return a JSON object following the specified format."
+    )
+
+    return EXTRACTOR_SYSTEM_PROMPT, "\n".join(user_prompt_parts)

kg_mcp/llm/prompts/linker.py

@@ -0,0 +1,117 @@
+"""
+Linker prompt template for entity deduplication and relationship inference.
+"""
+
+import json
+from typing import Any, Dict, List, Tuple
+
+from kg_mcp.llm.schemas import ExtractionResult
+
+
+LINKER_SYSTEM_PROMPT = """You are an expert at analyzing knowledge graphs and detecting duplicates/relationships.
+
+Your task is to:
+1. DETECT DUPLICATES: Check if any newly extracted entities are duplicates of existing ones
+2. SUGGEST MERGES: If entities are the same or very similar, suggest merging them
+3. INFER RELATIONSHIPS: Based on context, suggest relationships between entities
+
+RELATIONSHIP TYPES:
+- DECOMPOSES_INTO: A goal breaks down into subgoals
+- HAS_CONSTRAINT: A goal has a constraint
+- HAS_STRATEGY: A goal has a strategy for implementation
+- BLOCKED_BY: An entity is blocked by a pain point
+- IMPLEMENTED_BY: A goal is implemented by code artifacts
+- VERIFIED_BY: A goal is verified by tests
+- RELATED_TO: Generic relationship
+
+RULES:
+- Only suggest merges with high confidence (> 0.7)
+- Consider semantic similarity, not just exact matches
+- For relationships, provide clear reasoning
+- Use existing entity IDs when available
+
+OUTPUT FORMAT (JSON):
+{
+  "merge_suggestions": [
+    {
+      "new_entity_type": "Goal|Preference|etc",
+      "new_entity_title": "...",
+      "existing_entity_id": "uuid",
+      "existing_entity_title": "...",
+      "confidence": 0.85,
+      "reason": "Why they should be merged"
+    }
+  ],
+  "relationships": [
+    {
+      "source_type": "Goal",
+      "source_id": "uuid or null",
+      "source_title": "...",
+      "relationship_type": "DECOMPOSES_INTO|HAS_CONSTRAINT|etc",
+      "target_type": "SubGoal",
+      "target_id": "uuid or null",
+      "target_title": "...",
+      "confidence": 0.8
+    }
+  ]
+}"""
+
+
+def get_linker_prompt(
+    extraction: ExtractionResult,
+    existing_goals: List[Dict[str, Any]],
+    existing_preferences: List[Dict[str, Any]],
+    recent_interactions: List[Dict[str, Any]],
+) -> Tuple[str, str]:
+    """
+    Build the linker prompt for entity linking and relationship inference.
+
+    Args:
+        extraction: The extraction result to link
+        existing_goals: List of existing goals from the graph
+        existing_preferences: List of existing preferences
+        recent_interactions: Recent interactions for context
+
+    Returns:
+        Tuple of (system_prompt, user_prompt)
+    """
+    user_prompt_parts = []
+
+    # Add extracted entities
+    user_prompt_parts.append("NEWLY EXTRACTED ENTITIES:")
+    user_prompt_parts.append(f"```json\n{extraction.model_dump_json(indent=2)}\n```")
+
+    # Add existing goals
+    if existing_goals:
+        user_prompt_parts.append("\nEXISTING GOALS IN GRAPH:")
+        for goal in existing_goals[:20]:  # Limit to 20
+            user_prompt_parts.append(
+                f"- ID: {goal.get('id')}, Title: {goal.get('title')}, "
+                f"Status: {goal.get('status')}"
+            )
+
+    # Add existing preferences
+    if existing_preferences:
+        user_prompt_parts.append("\nEXISTING PREFERENCES:")
+        for pref in existing_preferences[:10]:
+            user_prompt_parts.append(
+                f"- ID: {pref.get('id')}, Category: {pref.get('category')}, "
+                f"Preference: {pref.get('preference')}"
+            )
+
+    # Add recent interaction context
+    if recent_interactions:
+        user_prompt_parts.append("\nRECENT INTERACTIONS (for context):")
+        for interaction in recent_interactions[:5]:
+            user_prompt_parts.append(
+                f"- {interaction.get('timestamp', 'N/A')}: "
+                f"{interaction.get('user_text', '')[:100]}..."
+            )
+
+    user_prompt_parts.append(
+        "\n\nAnalyze the newly extracted entities against the existing graph. "
+        "Suggest merges for duplicates and infer relationships. "
+        "Return a JSON object following the specified format."
+    )
+
+    return LINKER_SYSTEM_PROMPT, "\n".join(user_prompt_parts)

kg_mcp/llm/schemas.py

@@ -0,0 +1,248 @@
+"""
+Pydantic schemas for LLM input/output validation.
+These define the structured format for entity extraction and linking.
+"""
+
+from datetime import datetime
+from typing import List, Optional
+from uuid import uuid4
+
+from pydantic import BaseModel, Field
+
+
+# =============================================================================
+# Extraction Schemas
+# =============================================================================
+
+
+class GoalExtract(BaseModel):
+    """Extracted goal from user text."""
+
+    title: str = Field(..., description="Short title of the goal")
+    description: Optional[str] = Field(None, description="Detailed description")
+    priority: int = Field(default=2, ge=1, le=5, description="Priority 1-5 (1=highest)")
+    status: str = Field(default="active", description="Status: active, paused, done")
+    parent_goal_title: Optional[str] = Field(
+        None, description="Title of parent goal if this is a subgoal"
+    )
+
+
+class ConstraintExtract(BaseModel):
+    """Extracted constraint from user text."""
+
+    type: str = Field(..., description="Constraint type: budget, stack, style, performance, time")
+    description: str = Field(..., description="Description of the constraint")
+    severity: str = Field(default="must", description="Severity: must, should, nice_to_have")
+
+
+class PreferenceExtract(BaseModel):
+    """Extracted preference from user text."""
+
+    category: str = Field(
+        ...,
+        description="Category: coding_style, architecture, testing, tools, output_format",
+    )
+    preference: str = Field(..., description="The preference itself")
+    strength: str = Field(default="prefer", description="Strength: prefer, avoid, require")
+
+
+class PainPointExtract(BaseModel):
+    """Extracted pain point from user text."""
+
+    description: str = Field(..., description="Description of the pain point")
+    severity: str = Field(default="medium", description="Severity: low, medium, high, critical")
+    related_goal: Optional[str] = Field(None, description="Related goal title if any")
+
+
+class StrategyExtract(BaseModel):
+    """Extracted strategy from user text."""
+
+    title: str = Field(..., description="Short title of the strategy")
+    approach: str = Field(..., description="Description of the approach")
+    rationale: Optional[str] = Field(None, description="Why this strategy was chosen")
+    outcome: Optional[str] = Field(None, description="Outcome of the strategy: success, failure, pending")
+    outcome_reason: Optional[str] = Field(None, description="Reason for the outcome")
+    related_goal: Optional[str] = Field(None, description="Related goal title if any")
+
+
+class AcceptanceCriteriaExtract(BaseModel):
+    """Extracted acceptance criteria from user text."""
+
+    criterion: str = Field(..., description="The acceptance criterion")
+    related_goal: Optional[str] = Field(None, description="Related goal title")
+    testable: bool = Field(default=True, description="Whether it's testable")
+
+
+class CodeReference(BaseModel):
+    """Reference to code in the message."""
+
+    path: str = Field(..., description="File path")
+    symbol: Optional[str] = Field(None, description="Symbol name (function/class)")
+    start_line: Optional[int] = Field(None, description="Start line number")
+    end_line: Optional[int] = Field(None, description="End line number")
+    action: str = Field(
+        default="reference", description="Action: reference, create, modify, delete"
+    )
+
+
+class ExtractionResult(BaseModel):
+    """Complete result of entity extraction from user text."""
+
+    goals: List[GoalExtract] = Field(default_factory=list)
+    constraints: List[ConstraintExtract] = Field(default_factory=list)
+    preferences: List[PreferenceExtract] = Field(default_factory=list)
+    pain_points: List[PainPointExtract] = Field(default_factory=list)
+    strategies: List[StrategyExtract] = Field(default_factory=list)
+    acceptance_criteria: List[AcceptanceCriteriaExtract] = Field(default_factory=list)
+    code_references: List[CodeReference] = Field(default_factory=list)
+    next_actions: List[str] = Field(default_factory=list)
+    confidence: float = Field(default=0.8, ge=0.0, le=1.0)
+
+
+# =============================================================================
+# Linking Schemas
+# =============================================================================
+
+
+class MergeSuggestion(BaseModel):
+    """Suggestion to merge a new entity with an existing one."""
+
+    new_entity_type: str = Field(..., description="Type of the new entity")
+    new_entity_title: str = Field(..., description="Title of the new entity")
+    existing_entity_id: str = Field(..., description="ID of the existing entity to merge with")
+    existing_entity_title: str = Field(..., description="Title of existing entity")
+    confidence: float = Field(default=0.8, ge=0.0, le=1.0)
+    reason: str = Field(..., description="Why these should be merged")
+
+
+class RelationshipSuggestion(BaseModel):
+    """Suggestion for a new relationship between entities."""
+
+    source_type: str = Field(..., description="Type of source entity")
+    source_id: Optional[str] = Field(None, description="ID of source entity (if existing)")
+    source_title: str = Field(..., description="Title of source entity")
+    relationship_type: str = Field(..., description="Type of relationship (e.g., IMPLEMENTED_BY)")
+    target_type: str = Field(..., description="Type of target entity")
+    target_id: Optional[str] = Field(None, description="ID of target entity (if existing)")
+    target_title: str = Field(..., description="Title of target entity")
+    confidence: float = Field(default=0.8, ge=0.0, le=1.0)
+
+
+class LinkingResult(BaseModel):
+    """Result of entity linking analysis."""
+
+    merge_suggestions: List[MergeSuggestion] = Field(default_factory=list)
+    relationships: List[RelationshipSuggestion] = Field(default_factory=list)
+
+
+# =============================================================================
+# Graph Node Schemas (for Neo4j)
+# =============================================================================
+
+
+class BaseNode(BaseModel):
+    """Base class for all graph nodes."""
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+    updated_at: datetime = Field(default_factory=datetime.utcnow)
+
+
+class InteractionNode(BaseNode):
+    """Represents a user interaction/request."""
+
+    user_text: str
+    assistant_text: Optional[str] = None
+    tags: List[str] = Field(default_factory=list)
+    project_id: str
+
+
+class GoalNode(BaseNode):
+    """Represents a goal or objective."""
+
+    title: str
+    description: Optional[str] = None
+    status: str = "active"
+    priority: int = 2
+    project_id: str
+
+
+class ConstraintNode(BaseNode):
+    """Represents a constraint."""
+
+    type: str
+    description: str
+    severity: str = "must"
+    project_id: str
+
+
+class PreferenceNode(BaseNode):
+    """Represents a user preference."""
+
+    category: str
+    preference: str
+    strength: str = "prefer"
+    user_id: str
+
+
+class PainPointNode(BaseNode):
+    """Represents a pain point."""
+
+    description: str
+    severity: str = "medium"
+    resolved: bool = False
+    project_id: str
+
+
+class StrategyNode(BaseNode):
+    """Represents a strategy or approach."""
+
+    title: str
+    approach: str
+    rationale: Optional[str] = None
+    outcome: Optional[str] = None
+    outcome_reason: Optional[str] = None
+    project_id: str
+
+
+class DecisionNode(BaseNode):
+    """Represents an ADR-lite decision."""
+
+    title: str
+    decision: str
+    rationale: str
+    alternatives: List[str] = Field(default_factory=list)
+    project_id: str
+
+
+class CodeArtifactNode(BaseNode):
+    """Represents a code artifact (file, function, class, snippet)."""
+
+    path: str
+    language: Optional[str] = None
+    kind: str = "file"  # file, function, class, snippet
+    git_commit: Optional[str] = None
+    content_hash: Optional[str] = None
+    start_line: Optional[int] = None
+    end_line: Optional[int] = None
+    project_id: str
+
+
+class SymbolNode(BaseNode):
+    """Represents a code symbol (function, class, method)."""
+
+    fqn: str  # fully qualified name
+    name: str
+    kind: str  # function, class, method, variable
+    signature: Optional[str] = None
+    artifact_id: str
+
+
+class TestCaseNode(BaseNode):
+    """Represents a test case."""
+
+    name: str
+    kind: str = "unit"  # unit, integration, e2e, manual
+    path: Optional[str] = None
+    status: str = "pending"  # pending, passed, failed
+    project_id: str

kg_mcp/main.py

@@ -0,0 +1,195 @@
+"""
+Main entry point for the MCP-KG-Memory server.
+Supports both STDIO and Streamable HTTP transports for Antigravity compatibility.
+
+Usage:
+    # STDIO mode (for Antigravity command/args config)
+    python -m kg_mcp --transport stdio
+
+    # HTTP mode (for Antigravity serverUrl config or standalone)
+    python -m kg_mcp --transport http --host 127.0.0.1 --port 8000
+"""
+
+import argparse
+import logging
+import os
+import signal
+import sys
+from typing import Optional
+
+from mcp.server.fastmcp import FastMCP
+
+from kg_mcp.config import get_settings
+from kg_mcp.mcp.tools import register_tools
+from kg_mcp.mcp.resources import register_resources
+from kg_mcp.mcp.prompts import register_prompts
+
+
+def setup_logging(transport: str) -> logging.Logger:
+    """
+    Configure logging based on transport mode.
+
+    For STDIO: Log ONLY to stderr (stdout is reserved for MCP protocol)
+    For HTTP: Log to stdout
+    """
+    settings = get_settings()
+    log_level = getattr(logging, settings.log_level.upper(), logging.INFO)
+
+    # Clear any existing handlers
+    root_logger = logging.getLogger()
+    root_logger.handlers.clear()
+
+    if transport == "stdio":
+        # CRITICAL: In stdio mode, stdout is for MCP protocol ONLY
+        # All logs must go to stderr
+        handler = logging.StreamHandler(sys.stderr)
+    else:
+        handler = logging.StreamHandler(sys.stdout)
+
+    handler.setFormatter(
+        logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    )
+    root_logger.addHandler(handler)
+    root_logger.setLevel(log_level)
+
+    return logging.getLogger(__name__)
+
+
+def create_mcp_server(json_response: bool = True, stateless: bool = True) -> FastMCP:
+    """Create and configure the MCP server instance."""
+    logger = logging.getLogger(__name__)
+    logger.info("Initializing MCP-KG-Memory Server...")
+
+    # Create FastMCP instance
+    mcp = FastMCP(
+        "KG Memory Server",
+        json_response=json_response,
+        stateless_http=stateless,
+    )
+
+    # Register all MCP components
+    register_tools(mcp)
+    register_resources(mcp)
+    register_prompts(mcp)
+
+    logger.info("MCP server components registered successfully")
+    return mcp
+
+
+def handle_shutdown(signum, frame):
+    """Handle graceful shutdown on SIGTERM/SIGINT."""
+    logger = logging.getLogger(__name__)
+    logger.info(f"Received signal {signum}, shutting down gracefully...")
+    sys.exit(0)
+
+
+def run_stdio():
+    """Run server in STDIO mode (for Antigravity command/args config)."""
+    logger = setup_logging("stdio")
+    logger.info("Starting MCP-KG-Memory Server in STDIO mode")
+
+    # Register signal handlers for graceful shutdown
+    signal.signal(signal.SIGTERM, handle_shutdown)
+    signal.signal(signal.SIGINT, handle_shutdown)
+
+    settings = get_settings()
+    if settings.kg_mcp_token:
+        logger.info("Token authentication configured")
+    else:
+        logger.warning("No authentication token configured")
+
+    logger.info(f"LLM Model: {settings.llm_model}")
+    logger.info(f"Neo4j URI: {settings.neo4j_uri}")
+
+    # Create and run server
+    mcp = create_mcp_server(json_response=True, stateless=True)
+
+    # Run with stdio transport
+    mcp.run(transport="stdio")
+
+
+def run_http(host: str = "127.0.0.1", port: int = 8000, path: str = "/mcp"):
+    """Run server in HTTP mode (for Antigravity serverUrl config or standalone)."""
+    logger = setup_logging("http")
+    logger.info(f"Starting MCP-KG-Memory Server in HTTP mode on {host}:{port}{path}")
+
+    # Register signal handlers
+    signal.signal(signal.SIGTERM, handle_shutdown)
+    signal.signal(signal.SIGINT, handle_shutdown)
+
+    settings = get_settings()
+    if settings.kg_mcp_token:
+        logger.info("Bearer token authentication enabled")
+    else:
+        logger.warning("⚠️  No authentication token configured! Set KG_MCP_TOKEN in .env")
+
+    logger.info(f"LLM Model: {settings.llm_model}")
+    logger.info(f"Neo4j URI: {settings.neo4j_uri}")
+
+    # Create and run server
+    mcp = create_mcp_server(json_response=True, stateless=True)
+
+    # Run with streamable-http transport
+    mcp.run(
+        transport="streamable-http",
+        host=host,
+        port=port,
+    )
+
+
+def main():
+    """Main entry point with CLI argument parsing."""
+    parser = argparse.ArgumentParser(
+        description="MCP-KG-Memory Server - Knowledge Graph Memory for IDE Agents",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # STDIO mode (for Antigravity command config)
+  python -m kg_mcp --transport stdio
+
+  # HTTP mode (for Antigravity serverUrl config)
+  python -m kg_mcp --transport http --host 127.0.0.1 --port 8000
+
+  # Using console script
+  kg-mcp --transport stdio
+        """,
+    )
+
+    parser.add_argument(
+        "--transport",
+        "-t",
+        choices=["stdio", "http"],
+        default="http",
+        help="Transport mode: 'stdio' for command-based, 'http' for serverUrl-based (default: http)",
+    )
+
+    parser.add_argument(
+        "--host",
+        default="127.0.0.1",
+        help="Host to bind to in HTTP mode (default: 127.0.0.1)",
+    )
+
+    parser.add_argument(
+        "--port",
+        "-p",
+        type=int,
+        default=8000,
+        help="Port to listen on in HTTP mode (default: 8000)",
+    )
+
+    parser.add_argument(
+        "--path",
+        default="/mcp",
+        help="MCP endpoint path in HTTP mode (default: /mcp)",
+    )
+
+    args = parser.parse_args()
+
+    if args.transport == "stdio":
+        run_stdio()
+    else:
+        run_http(host=args.host, port=args.port, path=args.path)
+
+
+if __name__ == "__main__":
+    main()

kg_mcp/mcp/__init__.py

@@ -0,0 +1,3 @@
+"""
+MCP submodule for tools, resources, and prompts.
+"""