Merge pull request #13 from SecrinLabs/feat/store-pr-metadata-as-knowledge-node

Add pull request metadata ingestion and update related models and edges

Author: jenilsavani9

packages/app_store/github/webhook.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Dict, Any, Optional
 from packages.ingest.incremental_ingest import incremental_ingest
+from packages.ingest.pr_metadata import ingest_pull_request_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -29,6 +30,15 @@ def handle_pull_request(payload: Dict[str, Any]):
 
     logger.info(f"Processing PR action: {action}, merged: {merged}")
 
+    # First, ingest PR metadata for any PR action
+    repo_info = payload.get("repository", {})
+    try:
+        pr_data = ingest_pull_request_metadata(pull_request, repo_info)
+        logger.info(f"Ingested PR metadata: {pr_data}")
+    except Exception as e:
+        logger.error(f"Failed to ingest PR metadata: {e}", exc_info=True)
+
+    # Then handle merged PR actions
     if action == "closed" and merged:
         # PR was merged
         repo_info = payload.get("repository", {})
@@ -48,7 +58,7 @@ def handle_pull_request(payload: Dict[str, Any]):
                 "branch": base_branch
             }
 
-    return {"status": "ignored", "reason": "Not a merged PR"}
+    return {"status": "processed", "reason": "PR metadata ingested"}
 
 def handle_push(payload: Dict[str, Any]):
     """

packages/ingest/add_embeddings.py

@@ -102,7 +102,7 @@ def add_embeddings_to_all_nodes(
         provider: Embedding provider to use
     """
     if node_types is None:
-        node_types = ["Function", "Class", "File", "Doc", "Module", "Commit"]
+        node_types = ["Function", "Class", "File", "Doc", "Module", "Commit", "PullRequest"]
 
     print(f"Using {provider.value} embeddings")
     print(f"Embedding dimension: {settings.EMBEDDING_DIMENSION}")
@@ -158,7 +158,7 @@ def main():
     parser.add_argument(
         "--node-types",
         nargs="+",
-        choices=["Function", "Class", "File", "Doc", "Module", "Commit"],
+        choices=["Function", "Class", "File", "Doc", "Module", "Commit", "PullRequest"],
         help="Node types to process (default: all)"
     )
     parser.add_argument(

packages/ingest/edges.py

@@ -20,3 +20,6 @@ class Edge(Enum):
     # Issues / PRs / context
     REFERENCES = "REFERENCES"
     RELATED_TO = "RELATED_TO"
+    CREATED_BY = "CREATED_BY"
+    MERGED_TO = "MERGED_TO"
+    INCLUDES = "INCLUDES"

packages/ingest/full_ingest.py

@@ -44,7 +44,7 @@ def full_ingest(repo_path: str, branch: Optional[str] = None, max_commits: int =
     # We need to ensure vector indexes exist first (usually handled by migrations, but let's assume they exist)
     # add_embeddings_to_all_nodes handles all node types
     add_embeddings_to_all_nodes(
-        node_types=["Function", "Class", "File", "Doc", "Module", "Commit"],
+        node_types=["Function", "Class", "File", "Doc", "Module", "Commit", "PullRequest"],
         provider=EmbeddingProvider(settings.EMBEDDING_PROVIDER)
     )
 

packages/ingest/incremental_ingest.py

@@ -110,7 +110,7 @@ def incremental_ingest(repo_url: str, branch: Optional[str] = None):
         # 5. Update Embeddings
         print("Updating embeddings...")
         add_embeddings_to_all_nodes(
-            node_types=["Function", "Class", "File", "Doc", "Module", "Commit"],
+            node_types=["Function", "Class", "File", "Doc", "Module", "Commit", "PullRequest"],
             provider=EmbeddingProvider(settings.EMBEDDING_PROVIDER)
         )
 

packages/ingest/pr_metadata.py

@@ -0,0 +1,132 @@
+from datetime import datetime
+from typing import Dict, Any, Optional
+from packages.memory.memory import Memory
+from packages.ingest.edges import Edge
+from packages.parser.utils import extract_repo_info
+
+def ingest_pull_request_metadata(
+    pr_payload: Dict[str, Any],
+    repo_payload: Dict[str, Any],
+    memory: Optional[Memory] = None
+) -> Dict[str, Any]:
+    """
+    Ingest pull request metadata into the knowledge graph.
+    
+    Args:
+        pr_payload: The pull_request object from GitHub webhook
+        repo_payload: The repository object from GitHub webhook
+        memory: Optional Memory instance (creates new if not provided)
+    
+    Returns:
+        Dictionary with created node IDs and stats
+    """
+    if memory is None:
+        memory = Memory()
+    
+    # Extract repository info
+    repo_url = repo_payload.get("clone_url") or repo_payload.get("html_url") or ""
+    if not repo_url:
+        raise ValueError("Repository URL not found in payload")
+    
+    repo_info = extract_repo_info(repo_url)
+    repo_name = repo_info.get("name", "unknown")
+    
+    # Extract PR metadata
+    pr_number = pr_payload.get("number")
+    pr_title = pr_payload.get("title", "")
+    pr_body = pr_payload.get("body", "")
+    pr_state = pr_payload.get("state", "open")
+    pr_merged = pr_payload.get("merged", False)
+    pr_merged_at = pr_payload.get("merged_at")
+    
+    # Extract author info
+    pr_author_data = pr_payload.get("user", {})
+    pr_author = pr_author_data.get("login", "unknown")
+    pr_author_email = pr_author_data.get("email", f"{pr_author}@github.com")
+    
+    # Extract branch info
+    base_branch = pr_payload.get("base", {}).get("ref")
+    head_branch = pr_payload.get("head", {}).get("ref")
+    
+    # Create stable PR ID
+    pr_id = f"{repo_name}:pr:{pr_number}"
+    
+    # Build PR content for embedding
+    pr_content = f"""Pull Request #{pr_number}: {pr_title}
+
+      {pr_body}
+
+      Author: {pr_author}
+      State: {pr_state}
+      Merged: {pr_merged}
+      Base Branch: {base_branch}
+      Head Branch: {head_branch}
+      """
+    
+    # Upsert PR node
+    pr_node_id = memory.upsert_node(
+        "PullRequest",
+        match_props={"pr_number": pr_number, "repo_url": repo_url},
+        set_props={
+            "id": pr_id,
+            "pr_number": pr_number,
+            "title": pr_title,
+            "body": pr_body,
+            "content": pr_content,
+            "author": pr_author,
+            "repo_url": repo_url,
+            "state": pr_state,
+            "merged": pr_merged,
+            "merged_at": pr_merged_at,
+            "base_branch": base_branch,
+            "head_branch": head_branch,
+        }
+    )
+    
+    # Upsert repository node
+    repo_node_id = f"repo:{repo_name}"
+    memory.upsert_node(
+        "Repository",
+        match_props={"url": repo_url},
+        set_props={
+            "id": repo_node_id,
+            "url": repo_url,
+            "name": repo_name,
+            "owner": repo_info.get("owner", "unknown"),
+            "full_name": repo_info.get("full_name", "unknown"),
+            "content": repo_info.get("full_name", "unknown"),
+        }
+    )
+    
+    # Upsert author as Person node
+    author_id = f"person:{pr_author}"
+    memory.upsert_node(
+        "Person",
+        match_props={"email": pr_author_email},
+        set_props={
+            "id": author_id,
+            "name": pr_author,
+            "email": pr_author_email,
+            "content": f"{pr_author} ({pr_author_email})",
+        }
+    )
+    
+    # Create relationships
+    memory.link(pr_node_id, Edge.BELONGS_TO, repo_node_id)
+    memory.link(pr_node_id, Edge.CREATED_BY, author_id)
+    
+    # If PR is merged and includes commits, link to commits
+    # This would require additional commit SHAs from the payload
+    # For now, we'll just store the PR metadata
+    
+    return {
+        "pr_node_id": pr_node_id,
+        "pr_number": pr_number,
+        "title": pr_title,
+        "author": pr_author,
+        "state": pr_state,
+        "merged": pr_merged,
+    }
+
+
+__all__ = ["ingest_pull_request_metadata"]

packages/parser/core/graph_ingestion.py

@@ -16,6 +16,7 @@
     ModuleNode,
     TestNode,
     IssueNode,
+    PullRequestNode,
     Relationship,
 )
 
@@ -64,6 +65,7 @@ def _create_constraints(self):
             "CREATE CONSTRAINT IF NOT EXISTS FOR (m:Module) REQUIRE m.id IS UNIQUE",
             "CREATE CONSTRAINT IF NOT EXISTS FOR (t:Test) REQUIRE t.id IS UNIQUE",
             "CREATE CONSTRAINT IF NOT EXISTS FOR (i:Issue) REQUIRE i.id IS UNIQUE",
+            "CREATE CONSTRAINT IF NOT EXISTS FOR (pr:PullRequest) REQUIRE pr.id IS UNIQUE",
         ]
 
         for constraint in constraints:
@@ -209,6 +211,8 @@ def _infer_label_from_id(self, node_id: str) -> str:
             return "Test"
         elif ":issue:" in node_id:
             return "Issue"
+        elif ":pr:" in node_id:
+            return "PullRequest"
         else:
             # Default fallback
             return "Node"
@@ -269,14 +273,16 @@ def get_repository_stats(self, repo_name: str) -> Dict[str, int]:
         OPTIONAL MATCH (f)-[:HAS_DOC]->(d:Doc)
         OPTIONAL MATCH (commit:Commit)-[:TOUCHED]->(f)
         OPTIONAL MATCH (f)-[:IMPORTS]->(p:Package)
+        OPTIONAL MATCH (pr:PullRequest)-[:BELONGS_TO]->(r)
         RETURN 
             count(DISTINCT f) as files,
             count(DISTINCT c) as classes,
             count(DISTINCT fn) as functions,
             count(DISTINCT t) as tests,
             count(DISTINCT d) as docs,
             count(DISTINCT commit) as commits,
-            count(DISTINCT p) as packages
+            count(DISTINCT p) as packages,
+            count(DISTINCT pr) as pull_requests
         """
 
         result = self.client.run_query(query, {"repo_name": repo_name})
@@ -291,6 +297,7 @@ def get_repository_stats(self, repo_name: str) -> Dict[str, int]:
                 "docs": record["docs"],
                 "commits": record["commits"],
                 "packages": record["packages"],
+                "pull_requests": record["pull_requests"],
             }
 
         return {
@@ -301,6 +308,7 @@ def get_repository_stats(self, repo_name: str) -> Dict[str, int]:
             "docs": 0,
             "commits": 0,
             "packages": 0,
+            "pull_requests": 0,
         }
 
 

packages/parser/models/__init__.py

@@ -9,6 +9,7 @@
     TestNode,
     DocNode,
     IssueNode,
+    PullRequestNode,
     PackageNode,
 )
 from .relationships import (
@@ -28,6 +29,7 @@
     "TestNode",
     "DocNode",
     "IssueNode",
+    "PullRequestNode",
     "PackageNode",
     "Relationship",
     "RelationshipType",

packages/parser/models/nodes.py

@@ -15,6 +15,7 @@ class NodeType(str, Enum):
     TEST = "Test"
     DOC = "Doc"
     ISSUE = "Issue"
+    PULL_REQUEST = "PullRequest"
     PACKAGE = "Package"
 
 
@@ -102,6 +103,19 @@ class IssueNode(BaseNode):
     state: str  # e.g., 'open', 'closed'
 
 
+class PullRequestNode(BaseNode):
+    pr_number: int
+    title: str
+    body: str
+    author: str
+    repo_url: str
+    state: str  # e.g., 'open', 'closed', 'merged'
+    merged: bool = False
+    merged_at: Optional[datetime] = None
+    base_branch: Optional[str] = None
+    head_branch: Optional[str] = None
+
+
 class PackageNode(BaseNode):
     name: str
     version: str