Knowledge Graph module

ai-ml/knowledge_graph/.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+.pytest_cache/

ai-ml/knowledge_graph/README.md

@@ -0,0 +1,85 @@
+# Knowledge Graph — Team Lambda
+
+Connects related content a user has ingested, by reusing the vectors already
+produced by the embedding pipeline. Two layers:
+
+- **Document graph** — connects whole documents that are topically related.
+- **Topic graph** — finer connections between individual chunks, within and
+  across documents.
+
+## Where this fits
+
+```
+ai-ml/
+├── ingestion/
+├── embedding/
+├── quiz_generator/
+└── knowledge_graph/   ← this module
+```
+
+## Install
+
+From `ai-ml/`:
+
+```bash
+pip install -r knowledge_graph/requirements.txt
+```
+
+(Most of these are likely already installed via `embedding/requirements.txt`
+— this file just makes the module's own dependencies explicit.)
+
+## Run the API
+
+```bash
+cd ai-ml
+uvicorn knowledge_graph.app.api.graph_routes:app --reload --port 8005
+```
+
+Docs: http://127.0.0.1:8005/docs
+
+All endpoints require `Authorization: Bearer <jwt>`, same as quiz_generator.
+
+| Endpoint | Method | Description |
+|---|---|---|
+| `/health` | GET | No auth — basic liveness check |
+| `/graph` | GET | Returns the authenticated user's current graph (nodes + edges) |
+| `/graph/rebuild` | POST | Rebuilds the graph from the user's current embedded content |
+| `/graph` | DELETE | Deletes all graph edges for the authenticated user |
+
+## Run the tests
+
+```bash
+cd ai-ml
+pytest knowledge_graph/tests/ -v
+```
+
+14 tests, covering: utility math (vectorizer, similarity, labeling), graph
+building against a real Chroma collection, related vs. unrelated document
+detection, empty-user handling, and delete behavior.
+
+## How it works, briefly
+
+1. `document_graph_builder.py` groups a user's chunks by document, averages
+   each document's chunk vectors, and compares every pair using cosine
+   similarity (`utils/similarity.py`).
+2. Pairs above `SIMILARITY_THRESHOLD_DOCUMENT` (default 0.5, see `config.py`)
+   become edges, each labeled with shared keywords (`utils/topic_labeler.py`
+   — no LLM, no paid API).
+3. `topic_graph_builder.py` does the same at the individual-chunk level.
+4. `graph_service.py` orchestrates both builders and persists edges via
+   `storage/graph_store.py` — a JSON file living alongside the existing
+   shared ChromaDB data directory (no new database introduced).
+5. `api/graph_routes.py` exposes it all over HTTP, authenticated with the
+   same JWT pattern as quiz_generator (`quiz_generator.app.auth`).
+
+## Not included in this version
+
+- Frontend visualization (this module returns data, not a rendered graph)
+- LLM-based relationship explanations (labeling uses keyword overlap only,
+  to avoid requiring a paid API key)
+
+## Known limitation
+
+`storage/graph_store.py` uses a simple JSON file, not a proper database —
+fine for how this is used today (rebuilt on-demand per user), but not
+safe under heavy concurrent writes. Worth revisiting if usage grows.

ai-ml/knowledge_graph/app/api/graph_routes.py

@@ -0,0 +1,61 @@
+"""
+Knowledge Graph API — FastAPI service.
+
+Run from ai-ml/ (so the knowledge_graph.* and embedding.* imports resolve):
+  uvicorn knowledge_graph.app.api.graph_routes:app --reload --port 8003
+
+Interactive docs: http://127.0.0.1:8003/docs
+
+[Built in from the start, per NV-2] Every endpoint below requires a
+valid JWT, reusing quiz_generator.app.auth's existing, working
+implementation — the same fix already applied to ingestion. user_id
+is never accepted from the client; it comes only from the verified
+token's `sub` claim.
+"""
+from fastapi import FastAPI, Depends
+from fastapi.middleware.cors import CORSMiddleware
+
+from quiz_generator.app.auth import get_current_user_id
+from knowledge_graph.app.services.graph_service import GraphService
+
+app = FastAPI(title="StudyMind Knowledge Graph API — Team Lambda")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:5173"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+_service: GraphService | None = None
+
+
+def get_service() -> GraphService:
+    global _service
+    if _service is None:
+        _service = GraphService()
+    return _service
+
+
+@app.get("/health")
+def health_check():
+    return {"status": "ok"}
+
+
+@app.get("/graph")
+def get_graph_endpoint(user_id: str = Depends(get_current_user_id)) -> dict:
+    """Returns the authenticated user's current knowledge graph."""
+    return get_service().get_graph(user_id)
+
+
+@app.post("/graph/rebuild")
+def rebuild_graph_endpoint(user_id: str = Depends(get_current_user_id)) -> dict:
+    """Rebuilds the authenticated user's knowledge graph from their current content."""
+    return get_service().build_graph(user_id)
+
+
+@app.delete("/graph")
+def delete_graph_endpoint(user_id: str = Depends(get_current_user_id)) -> dict:
+    """Deletes all graph edges for the authenticated user."""
+    return get_service().delete_graph(user_id)

ai-ml/knowledge_graph/app/builders/document_graph_builder.py

@@ -0,0 +1,81 @@
+"""
+Builds document-to-document edges: groups a user's chunks by
+document, averages each document's chunk vectors, compares every
+pair of documents, and creates an edge wherever similarity is above
+the configured threshold.
+"""
+from itertools import combinations
+
+from embedding.chroma_store import get_collection
+from knowledge_graph.app.config import SIMILARITY_THRESHOLD_DOCUMENT, CONTENT_COLLECTION_NAME
+from knowledge_graph.app.utils.vectorizer import get_document_vector
+from knowledge_graph.app.utils.similarity import cosine_similarity
+from knowledge_graph.app.utils.topic_labeler import generate_label
+from knowledge_graph.app.validators.graph_validators import validate_documents_have_vectors
+from knowledge_graph.app.models.graph_edge import GraphEdge
+
+
+def get_user_documents(user_id: str) -> dict:
+    """
+    Fetches this user's chunks from the shared ChromaDB collection and
+    groups them by document_id.
+
+    Returns: { document_id: {"title": str, "vectors": [...], "texts": [...]} }
+    """
+    collection = get_collection(name=CONTENT_COLLECTION_NAME)
+    result = collection.get(
+        where={"user_id": user_id},
+        include=["embeddings", "documents", "metadatas"],
+    )
+
+    docs: dict = {}
+    ids = result.get("ids", [])
+    embeddings = result.get("embeddings", [])
+    documents = result.get("documents", [])
+    metadatas = result.get("metadatas", [])
+
+    for i in range(len(ids)):
+        meta = metadatas[i] or {}
+        doc_id = meta.get("document_id")
+        if not doc_id:
+            continue
+        docs.setdefault(doc_id, {"title": meta.get("document", ""), "vectors": [], "texts": []})
+        docs[doc_id]["vectors"].append(embeddings[i])
+        docs[doc_id]["texts"].append(documents[i])
+
+    return docs
+
+
+def build_document_graph(user_id: str) -> list:
+    """
+    Returns a list of GraphEdge.to_dict() for this user's
+    document-to-document relationships.
+    """
+    docs = validate_documents_have_vectors(get_user_documents(user_id))
+
+    doc_vectors = {
+        doc_id: get_document_vector(data["vectors"])
+        for doc_id, data in docs.items()
+    }
+
+    edges = []
+    for (id_a, vec_a), (id_b, vec_b) in combinations(doc_vectors.items(), 2):
+        score = cosine_similarity(vec_a, vec_b)
+        if score >= SIMILARITY_THRESHOLD_DOCUMENT:
+            # label from a sample of each document's text, not the full text
+            sample_a = " ".join(docs[id_a]["texts"][:2])
+            sample_b = " ".join(docs[id_b]["texts"][:2])
+
+            edge = GraphEdge(
+                user_id=user_id,
+                source_id=id_a,
+                target_id=id_b,
+                node_type="document",
+                similarity=round(score, 4),
+                source_title=docs[id_a]["title"],
+                target_title=docs[id_b]["title"],
+                label=generate_label(sample_a, sample_b),
+            )
+            edges.append(edge.to_dict())
+
+    return edges

ai-ml/knowledge_graph/app/builders/topic_graph_builder.py

@@ -0,0 +1,55 @@
+"""
+Builds finer-grained edges between individual chunks ("topics"),
+within and across a user's documents — the detail layer the
+document-level graph misses.
+"""
+from itertools import combinations
+
+from knowledge_graph.app.config import SIMILARITY_THRESHOLD_TOPIC
+from knowledge_graph.app.utils.similarity import cosine_similarity
+from knowledge_graph.app.utils.topic_labeler import generate_label
+from knowledge_graph.app.models.graph_edge import GraphEdge
+from knowledge_graph.app.builders.document_graph_builder import get_user_documents
+
+
+def build_topic_graph(user_id: str, max_chunks: int = 200) -> list:
+    """
+    Returns a list of GraphEdge.to_dict() for chunk-level relationships.
+
+    max_chunks caps how many of the user's chunks are compared, since
+    comparing every pair is O(n^2) — protects against runaway cost for
+    users with very large amounts of content. Chunks beyond the cap
+    are simply not included in this pass.
+    """
+    docs = get_user_documents(user_id)
+
+    # flatten into a single list of (chunk_id, vector, text, title)
+    chunks = []
+    for doc_id, data in docs.items():
+        for i, (vector, text) in enumerate(zip(data["vectors"], data["texts"])):
+            chunks.append({
+                "chunk_id": f"{doc_id}_{i}",
+                "vector": vector,
+                "text": text,
+                "title": data["title"],
+            })
+
+    chunks = chunks[:max_chunks]
+
+    edges = []
+    for a, b in combinations(chunks, 2):
+        score = cosine_similarity(a["vector"], b["vector"])
+        if score >= SIMILARITY_THRESHOLD_TOPIC:
+            edge = GraphEdge(
+                user_id=user_id,
+                source_id=a["chunk_id"],
+                target_id=b["chunk_id"],
+                node_type="topic",
+                similarity=round(score, 4),
+                source_title=a["title"],
+                target_title=b["title"],
+                label=generate_label(a["text"], b["text"]),
+            )
+            edges.append(edge.to_dict())
+
+    return edges

ai-ml/knowledge_graph/app/config.py

@@ -0,0 +1,36 @@
+"""
+Configuration for the knowledge graph module.
+
+Reuses the shared .env at the ai-ml/ root, same pattern as
+embedding/config.py and quiz_generator/app/auth.py.
+"""
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+
+# ai-ml/.env  (app -> knowledge-graph -> ai-ml)
+load_dotenv(dotenv_path=Path(__file__).resolve().parent.parent.parent / ".env")
+
+# --- Similarity thresholds ---
+# Document-level pairs below this score are not considered related.
+SIMILARITY_THRESHOLD_DOCUMENT = float(os.getenv("KG_DOC_SIMILARITY_THRESHOLD", "0.5"))
+
+# Topic/chunk-level pairs below this score are not considered related.
+# Slightly higher than the document threshold, since chunk-level
+# comparisons are noisier (less text per vector).
+SIMILARITY_THRESHOLD_TOPIC = float(os.getenv("KG_TOPIC_SIMILARITY_THRESHOLD", "0.6"))
+
+# --- Storage ---
+# Reuses the same local ChromaDB path as embedding/chroma_store.py,
+# so this module reads the same underlying data directory rather than
+# introducing a second database location. Edge data is written to its
+# own JSON file inside that same directory (see storage/graph_store.py) —
+# not a new database technology, just a new file alongside Chroma's data.
+CHROMA_DB_PATH = os.getenv(
+    "CHROMA_DB_PATH",
+    r"C:\Dev\QuantumLearningWorkspace\shared_chroma_data",
+)
+CONTENT_COLLECTION_NAME = "study_chunks"  # matches chroma_store.py's DEFAULT_COLLECTION_NAME
+
+# --- API ---
+API_PORT = int(os.getenv("KNOWLEDGE_GRAPH_PORT", "8005"))

ai-ml/knowledge_graph/app/models/graph_edge.py

@@ -0,0 +1,20 @@
+"""
+Represents a relationship (edge) between two nodes in the graph.
+"""
+from dataclasses import dataclass, asdict
+from typing import Optional
+
+
+@dataclass
+class GraphEdge:
+    user_id: str
+    source_id: str
+    target_id: str
+    node_type: str          # "document" or "topic" — matches the nodes it connects
+    similarity: float
+    source_title: str
+    target_title: str
+    label: Optional[str] = None   # short human-readable reason for the connection
+
+    def to_dict(self) -> dict:
+        return asdict(self)

ai-ml/knowledge_graph/app/models/graph_node.py

@@ -0,0 +1,16 @@
+"""
+Represents a single node in the knowledge graph — either a whole
+document, or (for the topic-level graph) an individual chunk.
+"""
+from dataclasses import dataclass, asdict
+
+
+@dataclass
+class GraphNode:
+    node_id: str          # document_id, or "documentid_chunkindex" for topic nodes
+    node_type: str         # "document" or "topic"
+    title: str
+    user_id: str
+
+    def to_dict(self) -> dict:
+        return asdict(self)