feat: add basic knowledge extraction and query

.gitignore

@@ -34,6 +34,8 @@ venv/
 env/
 python/.venv/
 .python-version
+llm_config.yaml
+python/knowledge_graph_data/
 
 # IDE & tooling
 .vscode/

README.md

@@ -5,7 +5,9 @@ Lance Graph is a Cypher-capable graph query engine built in Rust with Python bin
 This repository contains:
 
 - `rust/lance-graph` – the Cypher-capable query engine implemented in Rust
-- `python/` – PyO3 bindings and a thin `lance_graph` Python package
+- `python/` – PyO3 bindings and Python packages:
+  - `lance_graph` – thin wrapper around the Rust query engine
+  - `knowledge_graph` – Lance-backed knowledge graph CLI, API, and utilities
 
 ## Prerequisites
 
@@ -62,6 +64,65 @@ result = query.execute({"Person": people})
 print(result.to_pydict())  # {'name': ['Bob', 'David'], 'age': [34, 42]}
 ```
 
+## Knowledge Graph CLI & API
+
+The `knowledge_graph` package layers a simple Lance-backed knowledge graph
+service on top of the `lance_graph` engine. It provides:
+
+- A CLI (`knowledge_graph.main`) for initializing storage, running Cypher
+  queries, and bootstrapping data via heuristic text extraction.
+- A reusable FastAPI component, plus a standalone web service
+  (`knowledge_graph.webservice`) that exposes query and dataset endpoints.
+- Storage helpers that persist node and relationship tables as Lance datasets.
+
+### CLI usage
+
+```bash
+uv run knowledge_graph --init                    # initialize storage and schema stub
+uv run knowledge_graph --list-datasets           # list Lance datasets on disk
+uv run knowledge_graph --extract-preview notes.txt
+uv run knowledge_graph --extract-preview "Alice joined the graph team"
+uv run knowledge_graph --extract-and-add notes.txt
+uv run knowledge_graph "MATCH (n) RETURN n LIMIT 5"
+uv run knowledge_graph --log-level DEBUG --extract-preview "Inline text"
+uv run knowledge_graph --ask "Who is working on the Presto project?"
+
+
+# Configure LLM extraction (default)
+uv sync --extra llm  # install optional LLM dependencies
+uv sync --extra lance-storage  # install Lance dataset support
+export OPENAI_API_KEY=sk-...
+uv run knowledge_graph --llm-model gpt-4o-mini --extract-preview notes.txt
+
+# Supply additional OpenAI client options via YAML (base_url, headers, etc.)
+uv run knowledge_graph --llm-config llm_config.yaml --extract-and-add notes.txt
+
+# Fall back to the heuristic extractor when LLM access is unavailable
+uv run knowledge_graph --extractor heuristic --extract-preview notes.txt
+
+```
+
+The default extractor uses OpenAI. Configure credentials via environment
+variables supported by the SDK (for example `OPENAI_API_BASE` or
+`OPENAI_API_KEY`), or place them in a YAML file passed through `--llm-config`.
+Override the model and temperature with `--llm-model` and `--llm-temperature`.
+```
+
+By default the CLI writes datasets under `./knowledge_graph_data`. Provide
+`--root` and `--schema` to point at alternate storage locations and schema YAML.
+
+### FastAPI service
+
+Run the web service after installing the `knowledge_graph` package (and
+dependencies such as FastAPI):
+
+```bash
+uv run --package knowledge_graph knowledge_graph-webservice
+```
+
+The service exposes endpoints under `/graph`, including `/graph/health`,
+`/graph/query`, `/graph/datasets`, and `/graph/schema`.
+
 ### Development workflow
 
 For linting and type checks:

python/README.md

@@ -43,6 +43,14 @@ maturin develop
 
 - `python/src/` – PyO3 bridge that exposes graph APIs to Python
 - `python/python/lance_graph/` – pure-Python wrapper and `__init__`
+- `python/python/knowledge_graph/` – CLI, FastAPI, and extractor utilities built on Lance
 - `python/python/tests/` – graph-centric functional tests
 
 Refer to the repository root `README.md` for information about the Rust crate.
+
+> Run CLI commands through `uv run knowledge_graph ...`. The default uses an
+> LLM-backed extractor; install the LLM extra with `uv sync --extra llm` (or
+> `uv pip install -e '.[llm]'`) and configure `OPENAI_API_KEY`. Install
+> `uv sync --extra lance-storage` to enable Lance dataset persistence. Supply
+> extra options (e.g., `base_url`, HTTP headers) via `--llm-config`. Use
+> `--extractor heuristic` to avoid LLM calls during testing or offline work.

python/llm_config.example.yaml

@@ -0,0 +1,10 @@
+# Example LLM client configuration (public-safe template)
+# Copy this file to llm_config.yaml and fill in your values.
+api_key: YOUR_API_KEY_HERE
+# Optional: override base URL for self-hosted or compatible APIs (e.g., OpenAI-compatible gateways)
+base_url: https://api.openai.com/v1
+# Optional: additional HTTP headers to send with each request
+default_headers:
+  # openai-organization: YOUR_ORG_ID
+  # Authorization is derived from api_key by the client; do not duplicate here
+  # Custom-Header: value

python/pyproject.toml

@@ -1,7 +1,13 @@
 [project]
 name = "lance-graph"
 dynamic = ["version"]
-dependencies = ["pyarrow>=14"]
+dependencies = [
+    "pyarrow>=14",
+    "pyyaml>=6.0",
+    "fastapi>=0.104.0",
+    "uvicorn>=0.24.0",
+    "pydantic>=2.0.0",
+]
 description = "Python bindings for the lance-graph Cypher engine"
 authors = [{ name = "Lance Devs", email = "dev@lancedb.com" }]
 license = { file = "LICENSE" }
@@ -37,6 +43,8 @@ build-backend = "maturin"
 [project.optional-dependencies]
 tests = ["pytest", "pyarrow>=14", "pandas"]
 dev = ["ruff", "pyright"]
+llm = ["openai>=1.52.0"]
+lance-storage = ["lance>=0.17.0"]
 
 [project.scripts]
 knowledge_graph = "knowledge_graph.main:main"

python/python/knowledge_graph/__init__.py

@@ -13,6 +13,19 @@
 except ImportError:  # pragma: no cover - builder is available in normal installs.
     GraphConfigBuilder = object  # type: ignore[assignment]
 
+from .component import KnowledgeGraphComponent
+from .config import KnowledgeGraphConfig, build_graph_config_from_mapping
+from .extraction import (
+    DEFAULT_STRATEGY,
+    BaseExtractor,
+    get_extractor,
+    preview_extraction,
+)
+from .extractors import HeuristicExtractor, LLMExtractor
+from .service import LanceKnowledgeGraph, create_default_service
+from .store import LanceGraphStore
+from .webservice import create_app
+
 TableMapping = Mapping[str, pa.Table]
 
 
@@ -100,4 +113,20 @@ def build(self) -> KnowledgeGraph:
         return KnowledgeGraph(config, self._datasets)
 
 
-__all__ = ["KnowledgeGraph", "KnowledgeGraphBuilder"]
+__all__ = [
+    "KnowledgeGraph",
+    "KnowledgeGraphBuilder",
+    "KnowledgeGraphConfig",
+    "build_graph_config_from_mapping",
+    "LanceGraphStore",
+    "LanceKnowledgeGraph",
+    "create_default_service",
+    "KnowledgeGraphComponent",
+    "create_app",
+    "DEFAULT_STRATEGY",
+    "BaseExtractor",
+    "get_extractor",
+    "preview_extraction",
+    "HeuristicExtractor",
+    "LLMExtractor",
+]

python/python/knowledge_graph/component.py

@@ -0,0 +1,111 @@
+"""Reusable FastAPI component for the Lance knowledge graph service."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import pyarrow as pa
+import yaml
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+
+from .config import KnowledgeGraphConfig
+from .service import LanceKnowledgeGraph
+from .store import LanceGraphStore
+
+
+class QueryRequest(BaseModel):
+    query: str
+
+
+class QueryResponse(BaseModel):
+    rows: List[Dict[str, Any]]
+    row_count: int
+
+
+class DatasetUpsertRequest(BaseModel):
+    records: List[Dict[str, Any]]
+    merge: bool = True
+
+
+class KnowledgeGraphComponent:
+    """Bundle FastAPI routes that expose the Lance knowledge graph."""
+
+    def __init__(self, config: Optional[KnowledgeGraphConfig] = None):
+        self._config = config or KnowledgeGraphConfig.default()
+        self._service: Optional[LanceKnowledgeGraph] = None
+        self.router = APIRouter(tags=["knowledge-graph"])
+        self._setup_routes()
+
+    def _get_service(self) -> LanceKnowledgeGraph:
+        if self._service is None:
+            try:
+                self._service = _create_service(self._config)
+            except FileNotFoundError as exc:
+                raise HTTPException(status_code=500, detail=str(exc)) from exc
+        return self._service
+
+    def _setup_routes(self) -> None:
+        @self.router.get("/health")
+        async def health() -> Dict[str, str]:
+            return {"status": "healthy", "service": "lance-knowledge-graph"}
+
+        @self.router.get("/datasets")
+        async def list_datasets() -> Dict[str, List[str]]:
+            service = self._get_service()
+            names = list(service.dataset_names())
+            return {"datasets": names}
+
+        @self.router.get("/datasets/{name}")
+        async def get_dataset(name: str, limit: int = 100) -> Dict[str, Any]:
+            service = self._get_service()
+            if not service.has_dataset(name):
+                raise HTTPException(
+                    status_code=404, detail=f"Dataset '{name}' not found"
+                )
+
+            table = service.load_table(name)
+            rows = table.to_pylist()
+            if limit is not None:
+                rows = rows[:limit]
+            return {"name": name, "row_count": len(rows), "rows": rows}
+
+        @self.router.post("/datasets/{name}")
+        async def upsert_dataset(
+            name: str, request: DatasetUpsertRequest
+        ) -> Dict[str, Any]:
+            if not request.records:
+                raise HTTPException(status_code=400, detail="records cannot be empty")
+
+            table = pa.Table.from_pylist(request.records)
+            service = self._get_service()
+            service.upsert_table(name, table, merge=request.merge)
+            return {"status": "ok", "dataset": name, "row_count": table.num_rows}
+
+        @self.router.post("/query", response_model=QueryResponse)
+        async def execute_query(request: QueryRequest) -> QueryResponse:
+            service = self._get_service()
+            result = service.query(request.query)
+            rows = result.to_pylist()
+            return QueryResponse(rows=rows, row_count=len(rows))
+
+        @self.router.get("/schema")
+        async def get_schema() -> Dict[str, Any]:
+            schema_path = self._config.resolved_schema_path()
+            if not schema_path.exists():
+                raise HTTPException(status_code=404, detail="Schema file not found")
+            with schema_path.open("r", encoding="utf-8") as handle:
+                payload = yaml.safe_load(handle) or {}
+            return {"path": str(schema_path), "schema": payload}
+
+    def close(self) -> None:
+        """Release retained resources."""
+        self._service = None
+
+
+def _create_service(config: KnowledgeGraphConfig) -> LanceKnowledgeGraph:
+    graph_config = config.load_graph_config()
+    storage = LanceGraphStore(config)
+    service = LanceKnowledgeGraph(graph_config, storage=storage)
+    service.ensure_initialized()
+    return service