# Phase 1 — Core

> Get the package to a working state: scaffolding, both store implementations, tests, and a README.

## Tasks

- [ ] **Project scaffolding** — `pyproject.toml`, src layout, directory structure
- [ ] **`GrafeoGraphStore`** — knowledge graph storage (`graph_store.py`)
- [ ] **`GrafeoGraphVectorStore`** — vector + graph retrieval (`graph_vector_store.py`)
- [ ] **`graph_document.py`** — GraphDocument conversion helpers
- [ ] **`__init__.py`** — public exports
- [ ] **Test suite** — `test_graph_store.py`, `test_graph_vector_store.py`, `test_retriever.py`, `conftest.py`
- [ ] **README** — overview, install, usage examples, comparison table

---

## Repository Structure

```
grafeo-langchain/
├── pyproject.toml
├── README.md
├── src/
│   └── grafeo_langchain/
│       ├── __init__.py
│       ├── graph_store.py
│       ├── graph_vector_store.py
│       └── graph_document.py
└── tests/
    ├── conftest.py
    ├── test_graph_store.py
    ├── test_graph_vector_store.py
    └── test_retriever.py
```

## pyproject.toml

```toml
[project]
name = "grafeo-langchain"
version = "0.1.0"
description = "LangChain graph store and vector store backed by GrafeoDB embedded graph database"
readme = "README.md"
license = "Apache-2.0"
requires-python = ">=3.12"
authors = [{ name = "GrafeoDB" }]
keywords = ["langchain", "graph-store", "graph-rag", "grafeo", "knowledge-graph", "vector-search"]
classifiers = [
    "Development Status :: 3 - Alpha",
    "Intended Audience :: Developers",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
    "Topic :: Database :: Database Engines/Servers",
]

dependencies = [
    "grafeo>=0.4",
    "langchain-core>=0.3,<1",
]

[project.optional-dependencies]
retriever = ["langchain-graph-retriever>=0.3"]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/grafeo_langchain"]

# ── Ruff ──────────────────────────────────────────────────────────────
[tool.ruff]
target-version = "py312"
line-length = 100

[tool.ruff.lint]
select = ["E", "F", "W", "I", "UP", "B", "SIM", "RUF"]

# ── Ty ────────────────────────────────────────────────────────────────
[tool.ty]
python-version = "3.12"

# ── Pytest ────────────────────────────────────────────────────────────
[tool.pytest.ini_options]
testpaths = ["tests"]

# ── UV ────────────────────────────────────────────────────────────────
[tool.uv]
dev-dependencies = [
    "pytest>=8",
    "ruff>=0.11",
    "ty>=0.1",
    "langchain>=0.3",
    "langchain-graph-retriever>=0.3",
]
```

---

## Interface Contracts

### 1. `GraphStore` — Knowledge graph storage

LangChain's `GraphStore` (from `langchain_community.graphs.graph_store`) is the base class for graph databases that store knowledge triples extracted by LLMs.

```python
class GraphStore(ABC):
    """Abstract class for graph stores."""

    @property
    @abstractmethod
    def get_schema(self) -> str: ...

    @property
    @abstractmethod
    def get_structured_schema(self) -> dict[str, Any]: ...

    @abstractmethod
    def query(self, query: str, params: dict | None = None) -> list[dict[str, Any]]: ...

    @abstractmethod
    def refresh_schema(self) -> None: ...

    def add_graph_documents(
        self,
        graph_documents: list[GraphDocument],
        include_source: bool = False,
    ) -> None: ...
```

### 2. `GraphVectorStore` — Vector + graph hybrid retrieval

LangChain's `GraphVectorStore` (from `langchain_core.graph_vectorstores`) combines a vector store with explicit document-to-document links.

```python
class GraphVectorStore(VectorStore):
    # Inherits from VectorStore:
    #   add_documents, add_texts, similarity_search, similarity_search_by_vector

    def traversal_search(self, query: str, k: int = 4, depth: int = 1, **kwargs) -> list[Document]: ...
    def mmr_traversal_search(self, query: str, k: int = 4, depth: int = 2, fetch_k: int = 100, lambda_mult: float = 0.5, **kwargs) -> list[Document]: ...
    def as_retriever(self, **kwargs) -> GraphRetriever: ...
```

### 3. `GraphDocument` — LLM-extracted graph data

```python
class GraphDocument(BaseModel):
    nodes: list[Node]
    relationships: list[Relationship]
    source: Document

class Node(BaseModel):
    id: str
    type: str = "Node"
    properties: dict[str, Any] = {}

class Relationship(BaseModel):
    source: Node
    target: Node
    type: str
    properties: dict[str, Any] = {}
```

---

## Implementation Reference

### GrafeoGraphStore

```python
# src/grafeo_langchain/graph_store.py
"""GrafeoDB-backed GraphStore for LangChain."""

from __future__ import annotations

from typing import Any

import grafeo
from langchain_core.documents import Document
from langchain_core.documents.base import Node, Relationship, GraphDocument


class GrafeoGraphStore:
    """LangChain GraphStore backed by GrafeoDB.

    Stores knowledge graph triples extracted by LLMs as native
    Grafeo nodes and edges. Supports direct GQL/Cypher queries.

    Args:
        db_path: Path to persistent database file. None = in-memory.
    """

    def __init__(self, db_path: str | None = None) -> None:
        self._db_path = db_path
        if db_path:
            self._db = grafeo.GrafeoDB(db_path)
        else:
            self._db = grafeo.GrafeoDB()

        if not self._db.has_property_index("node_id"):
            self._db.create_property_index("node_id")

        self._schema: dict[str, Any] | None = None

    @property
    def client(self) -> grafeo.GrafeoDB:
        """Return the underlying GrafeoDB instance."""
        return self._db

    @property
    def get_schema(self) -> str:
        """Return the graph schema as a string for LLM prompts."""
        schema = self.get_structured_schema
        parts = []
        if schema.get("labels"):
            parts.append(f"Node labels: {', '.join(schema['labels'])}")
        if schema.get("edge_types"):
            parts.append(f"Relationship types: {', '.join(schema['edge_types'])}")
        if schema.get("property_keys"):
            parts.append(f"Properties: {', '.join(schema['property_keys'])}")
        return "\n".join(parts) if parts else "Empty graph"

    @property
    def get_structured_schema(self) -> dict[str, Any]:
        """Return the structured graph schema."""
        if self._schema is None:
            self.refresh_schema()
        return self._schema or {}

    def query(self, query: str, params: dict | None = None) -> list[dict[str, Any]]:
        """Execute a GQL/Cypher query and return results as dicts."""
        result = self._db.execute(query)
        return [row for row in result]

    def refresh_schema(self) -> None:
        """Refresh the cached schema from the database."""
        self._schema = self._db.schema()

    def add_graph_documents(
        self,
        graph_documents: list[GraphDocument],
        include_source: bool = False,
    ) -> None:
        """Add LLM-extracted graph documents to the store."""
        for graph_doc in graph_documents:
            source_gid = None

            if include_source and graph_doc.source:
                source_gid = self._upsert_source_node(graph_doc.source)

            node_id_map: dict[str, int] = {}
            for node in graph_doc.nodes:
                gid = self._upsert_node(node)
                node_id_map[node.id] = gid

                if source_gid is not None:
                    self._db.create_edge(
                        source_id=gid,
                        target_id=source_gid,
                        edge_type="MENTIONED_IN",
                    )

            for rel in graph_doc.relationships:
                src_gid = node_id_map.get(rel.source.id)
                tgt_gid = node_id_map.get(rel.target.id)

                if src_gid is None:
                    src_gid = self._upsert_node(rel.source)
                    node_id_map[rel.source.id] = src_gid
                if tgt_gid is None:
                    tgt_gid = self._upsert_node(rel.target)
                    node_id_map[rel.target.id] = tgt_gid

                edge_type = rel.type.replace(" ", "_").replace("-", "_").upper()
                self._db.create_edge(
                    source_id=src_gid,
                    target_id=tgt_gid,
                    edge_type=edge_type,
                    properties=rel.properties,
                )

        self._schema = None  # Invalidate cache

    def _upsert_node(self, node: Node) -> int:
        """Insert or find a graph node, returning the Grafeo node ID."""
        matches = self._db.find_nodes_by_property("node_id", node.id)
        if matches:
            for key, value in node.properties.items():
                self._db.set_node_property(matches[0], key, value)
            return matches[0]

        label = node.type.replace(" ", "_").replace("-", "_")
        if not label.isidentifier():
            label = "Node"

        props = {"node_id": node.id, **node.properties}
        created = self._db.create_node(labels=[label], properties=props)
        return created.id

    def _upsert_source_node(self, doc: Document) -> int:
        """Store a source document as a node."""
        doc_id = doc.metadata.get("id", str(hash(doc.page_content[:200])))
        matches = self._db.find_nodes_by_property("node_id", f"source_{doc_id}")
        if matches:
            return matches[0]

        props = {
            "node_id": f"source_{doc_id}",
            "text": doc.page_content[:5000],
            **{k: v for k, v in doc.metadata.items() if isinstance(v, (str, int, float, bool))},
        }
        created = self._db.create_node(labels=["SourceDocument"], properties=props)
        return created.id

    def close(self) -> None:
        """Close the database."""
        self._db.close()
```

### GrafeoGraphVectorStore

```python
# src/grafeo_langchain/graph_vector_store.py
"""GrafeoDB-backed GraphVectorStore for LangChain."""

from __future__ import annotations

from typing import Any, Iterable

import grafeo
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore


_LINK_KEY = "__graph_links__"


class GrafeoGraphVectorStore(VectorStore):
    """Combined vector + graph store backed by GrafeoDB.

    Documents are stored as nodes with embedding vectors for similarity
    search. Documents can be explicitly linked to each other via metadata,
    enabling graph-enhanced retrieval patterns.

    Args:
        embedding: LangChain Embeddings instance for encoding text.
        db_path: Path to persistent database. None = in-memory.
        embedding_dimensions: Vector dimensionality.
    """

    def __init__(
        self,
        embedding: Embeddings,
        db_path: str | None = None,
        embedding_dimensions: int = 1536,
    ) -> None:
        self._embedding = embedding
        self._db_path = db_path
        self._dims = embedding_dimensions

        if db_path:
            self._db = grafeo.GrafeoDB(db_path)
        else:
            self._db = grafeo.GrafeoDB()

        if not self._db.has_property_index("doc_id"):
            self._db.create_property_index("doc_id")

        self._index_created = False

    @property
    def embeddings(self) -> Embeddings:
        return self._embedding

    def _ensure_index(self) -> None:
        if not self._index_created:
            try:
                self._db.create_vector_index(
                    label="Document",
                    property="embedding",
                    dimensions=self._dims,
                    metric="cosine",
                )
            except Exception:
                pass
            self._index_created = True

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: list[dict] | None = None,
        ids: list[str] | None = None,
        **kwargs: Any,
    ) -> list[str]:
        """Add texts with embeddings and optional graph links."""
        texts_list = list(texts)
        metadatas = metadatas or [{} for _ in texts_list]
        ids = ids or [str(i) for i in range(len(texts_list))]

        embeddings = self._embedding.embed_documents(texts_list)
        self._ensure_index()

        created_ids: list[str] = []
        for text, meta, doc_id, emb in zip(texts_list, metadatas, ids, embeddings):
            links = meta.pop(_LINK_KEY, [])

            props: dict[str, Any] = {
                "doc_id": doc_id,
                "text": text,
                "embedding": emb,
            }
            for k, v in meta.items():
                if isinstance(v, (str, int, float, bool)):
                    props[k] = v

            node = self._db.create_node(labels=["Document"], properties=props)

            for link in links:
                target_matches = self._db.find_nodes_by_property("doc_id", link.get("target_id", ""))
                edge_type = link.get("type", "LINKS_TO")
                for target_gid in target_matches:
                    self._db.create_edge(
                        source_id=node.id,
                        target_id=target_gid,
                        edge_type=edge_type,
                        properties=link.get("properties", {}),
                    )

            created_ids.append(doc_id)

        return created_ids

    def similarity_search_by_vector(
        self,
        embedding: list[float],
        k: int = 4,
        **kwargs: Any,
    ) -> list[Document]:
        """Find k most similar documents by vector."""
        self._ensure_index()

        results = self._db.vector_search(
            label="Document",
            property="embedding",
            query=embedding,
            k=k,
        )

        docs = []
        for node_id, distance in results:
            node = self._db.get_node(node_id)
            if node is None:
                continue
            props = node.properties()
            text = props.pop("text", "")
            props.pop("embedding", None)
            doc_id = props.pop("doc_id", str(node_id))
            docs.append(Document(
                page_content=text,
                metadata={"id": doc_id, "score": 1.0 - distance, **props},
            ))
        return docs

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        **kwargs: Any,
    ) -> list[Document]:
        """Embed query text and find similar documents."""
        embedding = self._embedding.embed_query(query)
        return self.similarity_search_by_vector(embedding, k=k, **kwargs)

    def traversal_search(
        self,
        query: str,
        k: int = 4,
        depth: int = 1,
        **kwargs: Any,
    ) -> list[Document]:
        """Vector search + graph traversal."""
        seeds = self.similarity_search(query, k=k)
        if not seeds:
            return seeds

        seen_ids: set[str] = {doc.metadata.get("id", "") for doc in seeds}
        traversed: list[Document] = list(seeds)

        for doc in seeds:
            doc_id = doc.metadata.get("id", "")
            if not doc_id:
                continue

            result = self._db.execute(
                f"MATCH (src {{doc_id: '{doc_id}'}})-[*1..{depth}]->(neighbor:Document) "
                f"RETURN neighbor LIMIT {k * 2}"
            )
            for node in result.nodes():
                props = node.properties()
                nid = props.get("doc_id", str(node.id))
                if nid in seen_ids:
                    continue
                seen_ids.add(nid)

                text = props.pop("text", "")
                props.pop("embedding", None)
                props.pop("doc_id", None)
                traversed.append(Document(
                    page_content=text,
                    metadata={"id": nid, "source": "graph_traversal", **props},
                ))

        return traversed[:k * 2]

    def mmr_traversal_search(
        self,
        query: str,
        k: int = 4,
        depth: int = 2,
        fetch_k: int = 100,
        lambda_mult: float = 0.5,
        **kwargs: Any,
    ) -> list[Document]:
        """MMR-diversified graph traversal.

        Uses Grafeo's native MMR search (SIMD-accelerated Rust) —
        no numpy or Python-side vector math needed.

        NOTE: Depends on grafeo.GrafeoDB.mmr_search() being available.
        See: H:\\GrafeoDB\\grafeo\\.claude\\todo\\4_language_bindings\\mmr-search.md
        """
        # Phase 1: Embed query
        query_embedding = self._embedding.embed_query(query)

        # Phase 2: Native MMR search for diverse seed documents
        self._ensure_index()
        mmr_results = self._db.mmr_search(
            label="Document",
            property="embedding",
            query=query_embedding,
            k=k,
            fetch_k=fetch_k,
            lambda_mult=lambda_mult,
        )

        seeds = self._results_to_documents(mmr_results)
        if depth < 1 or not seeds:
            return seeds

        # Phase 3: Graph traversal from MMR-selected seeds
        seen_ids: set[str] = {doc.metadata.get("id", "") for doc in seeds}
        traversed: list[Document] = list(seeds)

        for doc in seeds:
            doc_id = doc.metadata.get("id", "")
            if not doc_id:
                continue

            result = self._db.execute(
                f"MATCH (src {{doc_id: '{doc_id}'}})-[*1..{depth}]->(neighbor:Document) "
                f"RETURN neighbor LIMIT {k * 2}"
            )
            for node in result.nodes():
                props = node.properties()
                nid = props.get("doc_id", str(node.id))
                if nid in seen_ids:
                    continue
                seen_ids.add(nid)

                text = props.pop("text", "")
                props.pop("embedding", None)
                props.pop("doc_id", None)
                traversed.append(Document(
                    page_content=text,
                    metadata={"id": nid, "source": "graph_traversal", **props},
                ))

        return traversed[:k * 2]

    @classmethod
    def from_texts(
        cls,
        texts: list[str],
        embedding: Embeddings,
        metadatas: list[dict] | None = None,
        db_path: str | None = None,
        **kwargs: Any,
    ) -> GrafeoGraphVectorStore:
        """Create a store from texts."""
        store = cls(embedding=embedding, db_path=db_path, **kwargs)
        store.add_texts(texts, metadatas=metadatas)
        return store

    @classmethod
    def from_documents(
        cls,
        documents: list[Document],
        embedding: Embeddings,
        db_path: str | None = None,
        **kwargs: Any,
    ) -> GrafeoGraphVectorStore:
        """Create a store from LangChain Documents."""
        texts = [doc.page_content for doc in documents]
        metadatas = [doc.metadata for doc in documents]
        return cls.from_texts(texts, embedding, metadatas=metadatas, db_path=db_path, **kwargs)

    def _results_to_documents(self, results: list[tuple[int, float]]) -> list[Document]:
        """Convert vector_search/mmr_search results to LangChain Documents."""
        docs = []
        for node_id, distance in results:
            node = self._db.get_node(node_id)
            if node is None:
                continue
            props = node.properties()
            text = props.pop("text", "")
            props.pop("embedding", None)
            doc_id = props.pop("doc_id", str(node_id))
            docs.append(Document(
                page_content=text,
                metadata={"id": doc_id, "score": 1.0 - distance, **props},
            ))
        return docs

    def close(self) -> None:
        """Close the database."""
        self._db.close()
```

### `__init__.py`

```python
# src/grafeo_langchain/__init__.py
from grafeo_langchain.graph_store import GrafeoGraphStore
from grafeo_langchain.graph_vector_store import GrafeoGraphVectorStore

__all__ = ["GrafeoGraphStore", "GrafeoGraphVectorStore"]
```

---

## Data Model Mapping

| LangChain concept | Grafeo representation |
|---|---|
| `Node(id="alice", type="Person")` | `(:Person {node_id: "alice"})` |
| `Relationship(source, target, type="KNOWS")` | `(alice)-[:KNOWS]->(bob)` |
| `GraphDocument.source` | `(:SourceDocument {text: "...", node_id: "source_..."})` |
| `Document` (in GraphVectorStore) | `(:Document {doc_id: "...", text: "...", embedding: [...]})` |
| `__graph_links__` metadata | Grafeo edges between Document nodes |
| `store.query("MATCH ...")` | `db.execute("MATCH ...")` |
| `traversal_search(depth=2)` | `MATCH (n)-[*1..2]->(m)` variable-length path |
| Vector similarity search | `db.vector_search(label, property, query, k)` |

## README Content

The README should include:

- Overview / elevator pitch
- Install instructions (`pip install grafeo-langchain`)
- Quick-start examples (GraphStore + GraphVectorStore)
- Comparison table vs Neo4j

### Usage Examples

#### Basic GraphStore

```python
from langchain_openai import ChatOpenAI
from langchain_experimental.graph_transformers import LLMGraphTransformer
from langchain_core.documents import Document
from grafeo_langchain import GrafeoGraphStore

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
transformer = LLMGraphTransformer(llm=llm)

documents = [
    Document(page_content="Alice works at Microsoft. Bob works at Google. Alice knows Bob."),
]
graph_documents = transformer.convert_to_graph_documents(documents)

store = GrafeoGraphStore(db_path="./knowledge.db")
store.add_graph_documents(graph_documents, include_source=True)

results = store.query("MATCH (p:Person)-[:WORKS_AT]->(c) RETURN p.node_id, c.node_id")
print(store.get_schema)
```

#### GraphVectorStore with Traversal

```python
from langchain_openai import OpenAIEmbeddings
from grafeo_langchain import GrafeoGraphVectorStore

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
store = GrafeoGraphVectorStore(
    embedding=embeddings,
    db_path="./doc_graph.db",
    embedding_dimensions=1536,
)

store.add_texts(
    texts=["Python is a programming language...", "Guido van Rossum...", "ABC influenced..."],
    metadatas=[
        {"id": "python", "__graph_links__": [{"target_id": "abc", "type": "INFLUENCED_BY"}]},
        {"id": "guido"},
        {"id": "abc", "__graph_links__": [{"target_id": "python", "type": "INFLUENCED"}]},
    ],
    ids=["python", "guido", "abc"],
)

docs = store.similarity_search("What programming languages exist?", k=2)
docs = store.traversal_search("What programming languages exist?", k=4, depth=2)
docs = store.mmr_traversal_search("programming history", k=4, depth=2, lambda_mult=0.7)
```

## Key Differentiators

| Feature | Neo4j GraphStore | Grafeo GraphStore |
|---|---|---|
| Requires server | Yes (Docker/Cloud) | **No** (embedded, pip install) |
| GraphStore | Yes | **Yes** |
| GraphVectorStore | Community package | **Built-in** (native HNSW) |
| Query language | Cypher | **GQL + Cypher + Gremlin** |
| Graph algorithms | GDS plugin ($$$) | **Built-in** (PageRank, Louvain, ...) |
| Deployment | Docker container | **Single .db file** |
| Offline/edge | No | **Yes** |
