PyWrapAI-Graph — Library Reference

Single responsibility: Assemble the complete message list and call the LLM.

PyWrapAI-Graph sits between your application and PyWrapAI. It handles everything that needs to happen before the LLM call: loading the system prompt, injecting conversation history in the right position, optionally retrieving and injecting RAG context, and packaging the current user message. The LLM never sees individual components — it only ever receives a single, fully assembled list of messages.

Phase 2 adds the full RAG pipeline: Embedder (in PyWrapAI) produces vectors, InMemoryVectorStore stores and searches them, and VectorRetriever ties them together. GraphLLM(retriever=...) activates automatic RAG context injection on every chat() call.

---

Why This Separation Exists

Without PyWrapAI-Graph, your application would need to build this structure manually before every call:

[System Prompt]     ← loaded from a file or database
[History Turn 1]    ← previous user message
[History Turn 1]    ← previous assistant response
[History Turn N…]   ← all past turns up to max_turns
[Current Message]   ← what the user just typed

Getting this order wrong produces incoherent or broken conversations. PyWrapAI-Graph enforces the correct order every time, and PromptRegistry makes system prompt management a matter of convention rather than code.

---

Public API

from pywrapai_graph import (
    # Core
    GraphLLM, PromptRegistry,
    # Tools
    Tool, ToolRegistry, tool,
    # Agent (Phase 3)
    Agent, AgentConfig, AgentStep, AgentResult, AgentMaxTurnsError,
    # RAG pipeline — data types + chunking
    Document, Chunk, ScoredChunk,
    Chunker,
    # Vector stores — all backends
    BaseVectorStore,
    InMemoryVectorStore,          # built-in, no extra deps
    SQLiteVectorStore,            # stdlib sqlite3, no extra deps
    ChromaVectorStore,            # pip install chromadb
    FAISSVectorStore,             # pip install faiss-cpu
    PGVectorStore,                # pip install psycopg2-binary pgvector
    QdrantVectorStore,            # pip install qdrant-client
    # Retrieval
    BaseRetriever, VectorRetriever,
    retrieval_tool,
)

---

Directory Convention

Prompt files live in a directory you control. Pass the path explicitly via prompts_dir= so the library never depends on the process’s current working directory. The recommended layout:

my-project/
├── system_prompts_registry/
│   ├── support-agent-v1.txt
│   ├── support-agent-v2.txt
│   └── finance-bot-v1.txt
├── main.py
└── pywrapai_library.db

Prompt files are plain .txt files containing the system instruction. There are no special formatting rules — write the prompt exactly as you would pass it to the API.

Example — support-agent-v1.txt:

You are a friendly and professional customer support agent for Acme Corp.

Help users with:
- Product questions and specifications
- Account and billing issues
- Technical troubleshooting

Always be concise, empathetic, and solution-focused.
Never discuss pricing with competitors.
If you cannot help, offer to escalate to a human agent.

Naming convention (recommended): {role}-v{version}.txt

This makes it easy to run A/B tests between prompt versions and to track which version a session used.

---

GraphLLM — Constructor

from pywrapai       import LLM
from pywrapai_graph import GraphLLM

llm = LLM(provider="anthropic", model="claude-haiku-4-5-20251001")

graph = GraphLLM(
    llm         = llm,                             # required
    enable_system_prompt_versioning = True,         # default: False
    prompt_file = "support-agent-v1.txt",          # required if versioning=True
    prompts_dir = "/abs/path/to/system_prompts_registry",  # explicit path
    retriever   = my_retriever,                    # Phase 2: BaseRetriever instance
    rag_k       = 4,                               # Phase 2: number of chunks to retrieve
)

Parameters:

Parameter	Type	Default	Description
llm	LLM	required	The PyWrapAI LLM instance to call
enable_system_prompt_versioning	bool	False	Load a system prompt from the prompts directory
prompt_file	str	None	Filename inside the prompts directory. Required when versioning is enabled
prompts_dir	str	"system_prompts_registry"	Path to the prompts directory. Pass an absolute path in production — never rely on CWD
retriever	BaseRetriever | None	None	Phase 2: if set, enables automatic RAG context injection before every chat() call
rag_k	int	4	Phase 2: maximum number of retrieved chunks to inject as context

If enable_system_prompt_versioning=True but prompt_file is not provided, a ValueError is raised immediately at construction time.

If enable_system_prompt_versioning=False, no system prompt is loaded and injected. You can still set a system prompt on the underlying LLM instance using LLM(system_prompt=...).

If retriever is None (default), all RAG-related parameters (rag=, rag_k) are ignored and chat() behaves identically to pre-Phase 2.

---

GraphLLM.chat()

response = graph.chat(input, history=None, tools=None, tool_choice=None, rag=True)

Parameters:

• input: str | list[Message] — the current user message as a string, or a list of messages if you want to pass multiple messages at once
• history: list[Message] | None — past conversation turns to inject, usually loaded from your own persistence layer; pass None or [] for a fresh conversation
• tools: list[dict] | None — provider-neutral tool definitions from ToolRegistry.to_schema(). Forwarded to LLM.chat() unchanged.
• tool_choice: str | None — "auto", "required", "none", or a specific tool name. Forwarded to LLM.chat() unchanged.
• rag: bool — default True. Controls whether RAG context is injected when a retriever was set at construction time. Pass rag=False to skip retrieval for a specific call (e.g. factual greeting, administrative message).

Returns: LLMResponse (same as LLM.chat()). When the model calls a tool, response.content is "" and response.tool_calls is a non-empty list.

Message assembly order:

1. System prompt (if enable_system_prompt_versioning=True — always position 0)
2. History turns (in chronological order: oldest first)
3. Current user message — optionally RAG-augmented if retriever is set and rag=True

If enable_system_prompt_versioning=True and the caller’s input list starts with a SYSTEM message, that system message is silently replaced by the loaded prompt file. This prevents accidental double system prompts.

RAG injection (only when retriever is set and rag=True): The current user message is used as the retrieval query. Retrieved chunks are formatted into the D4 template and the current user message is replaced with the augmented version before assembly. The original message text is what should be saved to history — the <context> block is never persisted.

---

GraphLLM.stream()

for chunk in graph.stream(input, history=None):
    print(chunk, end="", flush=True)

Identical message assembly as chat(). Delegates to LLM.stream().

Does not accept tools= — passing it raises ValueError. Agent loops that need tool calling must use chat(). Streaming with interleaved tool-use deltas is deferred to Phase 3.

Token usage is tracked: after the generator exhausts, the provider reads token counts from the final chunk’s metadata and records them in the TokenTracker.

---

GraphLLM.chat_structured()

result = graph.chat_structured(input, schema=MyModel, history=None)

Identical to LLM.chat_structured() but with history and system prompt injection.

---

GraphLLM.achat() and GraphLLM.astream()

Async versions of chat() and stream(). Same assembly logic, same token tracking behaviour.

achat() accepts tools= and tool_choice=. astream() does not accept tools= (raises ValueError).

response = await graph.achat("Hello", history=history)

# With tools
response = await graph.achat(
    "What's the weather in Paris?",
    history     = history,
    tools       = registry.to_schema(),
    tool_choice = "auto",
)

async for chunk in graph.astream("Write a story", history=history):
    print(chunk, end="")

---

Proxy Properties

GraphLLM exposes the underlying LLM properties directly:

graph.tokens        # TokenTracker — same as graph.llm.tokens
graph.cache         # ResponseCache | None
graph.model         # str — e.g. "claude-haiku-4-5-20251001"
graph.provider      # str — e.g. "anthropic"
graph.system_prompt # str | None — the loaded prompt text
graph.llm           # the underlying LLM instance

---

GraphLLM.with_system_prompt()

g = graph.with_system_prompt("You are a finance expert.")

Returns a shallow copy of this GraphLLM with a new system prompt set. The copy shares the underlying _llm instance — same cache, same TokenTracker — with the original.

This is the correct pattern for injecting a per-user or per-session system prompt at runtime without mutating the shared graph object.

Parameters:

• text: str — the system prompt text to set on the copy

Returns: A new GraphLLM instance with system_prompt == text. The original graph is unchanged.

Key properties of the copy:

Property	Behaviour
_llm instance	Shared — same in-memory cache; a cached response from one copy is reused by all copies
TokenTracker	Shared — all per-user copies contribute to the same cost total
_system_prompt	Set to text on the copy only
_versioning	Set to True on the copy so the prompt is injected at message-assembly time
_retriever	Preserved from original — the copy shares the same BaseRetriever and therefore the same VectorStore. RAG retrieval is per-query, not per-user, so sharing the store is correct
_rag_k	Preserved from original

Primary use case — serving multiple roles from one base graph:

base_graph    = GraphLLM(llm=llm)
finance_graph = base_graph.with_system_prompt("You are a finance expert.")
support_graph = base_graph.with_system_prompt("You are a support agent.")

# All three share the same cache and TokenTracker
r1 = finance_graph.chat("What is the P/E ratio?")
r2 = support_graph.chat("How do I reset my password?")

---

PromptRegistry

Used internally by GraphLLM. You can also use it directly to inspect or list available prompts.

from pathlib import Path
from pywrapai_graph import PromptRegistry

# Always pass an explicit absolute path
registry = PromptRegistry(prompts_dir=str(Path(__file__).parent / "system_prompts_registry"))

# Load a specific prompt file
text = registry.load("support-agent-v1.txt")
print(text)

# List all available files
available = registry.list_files()
print(available)   # ["finance-bot-v1.txt", "support-agent-v1.txt", "support-agent-v2.txt"]

registry.load(filename)

text = registry.load("finance-bot-v1.txt")

• Returns: The full file content as a string, stripped of leading and trailing whitespace
• Raises: FileNotFoundError if the system_prompts_registry/ directory does not exist, or if the file does not exist. The error message includes the list of available files.

registry.list_files()

files = registry.list_files()

• Returns: A sorted list of filenames in the prompts directory
• If the directory does not exist, returns an empty list

---

Tool System

Phase 1.2 adds a complete tool-calling layer to PyWrapAI-Graph. You define Python functions, decorate them with @tool, register them in a ToolRegistry, and the registry handles everything else: building the JSON Schema automatically, generating provider-specific formats, and executing tool calls returned by the LLM.

Import paths:

from pywrapai_graph import Tool, ToolRegistry, tool
# or, directly from the submodule:
from pywrapai_graph.tools import Tool, ToolRegistry, tool

---

@tool — Decorator

The @tool decorator inspects a function’s type hints and docstring at decoration time and builds a Tool object. No JSON Schema is written by hand.

from pywrapai_graph import tool

@tool
def get_weather(city: str, units: str = "celsius") -> str:
    """Get the current weather for a city."""
    return f"16°C, cloudy in {city}"

What the decorator does automatically:

1. Reads the first line of the docstring as the tool description
2. Reads each parameter’s type hint and maps it to a JSON Schema type
3. Marks parameters as required unless they are Optional[X] / X | None or have a default value (units has a default, so it is not required)
4. Builds a Tool object and attaches it to the function as wrapper._pywrapai_tool

The decorated function still works exactly like the original — you can call it directly:

result = get_weather(city="London")   # still works

Type hint → JSON Schema type mapping:

Python type	JSON Schema type
str	"string"
int	"integer"
float	"number"
bool	"boolean"
list	"array"
dict	"object"
No hint / unknown	"string" (safe default)

Optional parameters — both syntax forms supported:

from typing import Optional

@tool
def search(query: str, limit: Optional[int] = None) -> list:
    """Search the knowledge base."""
    ...

# Works identically with Python 3.10+ pipe syntax:
@tool
def search(query: str, limit: int | None = None) -> list:
    """Search the knowledge base."""
    ...

# query  → required (str, no default)
# limit  → NOT required (Optional/None-union, has default)

---

Tool — Dataclass

Tool is the internal representation of a callable function with its schema metadata attached.

from pywrapai_graph import Tool

t = Tool(
    name        = "get_weather",
    description = "Get the current weather for a city.",
    parameters  = {
        "type": "object",
        "properties": {
            "city":  {"type": "string"},
            "units": {"type": "string"},
        },
        "required": ["city"],
    },
    fn = get_weather,
)

Fields:

Field	Type	Description
name	str	The function name the LLM uses to identify this tool
description	str	One-line description shown to the LLM in the schema
parameters	dict	JSON Schema object describing the function’s arguments
fn	Callable	The actual Python callable to execute

Tool is callable directly — t("London") delegates to t.fn("London").

In practice, use the @tool decorator rather than constructing Tool manually. The decorator reads type hints correctly and handles edge cases (missing hints, Optional, defaults).

---

ToolRegistry — Registry

ToolRegistry stores Tool objects, generates provider-specific schemas, and executes tool calls returned by the LLM.

from pywrapai_graph import ToolRegistry

registry = ToolRegistry()

register(item)

registry.register(get_weather)    # @tool-decorated function — recommended
registry.register(my_tool_obj)    # Tool instance directly

Accepts either a @tool-decorated function or a Tool instance. Plain functions without the @tool decorator raise TypeError immediately, so mistakes are caught at startup rather than at runtime.

get(name)

t = registry.get("get_weather")   # returns Tool; raises KeyError if not found

execute(tool_call)

from pywrapai import ToolCall

tc     = ToolCall(id="call_abc", name="get_weather", arguments={"city": "London"})
result = registry.execute(tc)
# ToolResult(tool_call_id="call_abc", name="get_weather", content="16°C, cloudy in London")

Runs fn(**tool_call.arguments) and wraps the return value in a ToolResult.

If the function raises an exception, the exception message is returned as ToolResult.content instead of crashing:

# If get_weather raises ValueError("city not found"):
# ToolResult(tool_call_id="call_abc", name="get_weather", content="Error: city not found")

This gives the model a chance to recover — retry with different arguments, ask for clarification, or call a different tool.

to_schema() — Primary method (Phase 1.6)

Returns all tools in the provider-neutral format. This is what you always pass to LLM.chat(tools=...) and GraphLLM.chat(tools=...). Each provider converts this internally to its own wire format — callers never pick a provider-specific schema method.

schema = registry.to_schema()
# [
#   {
#     "name":        "get_weather",
#     "description": "Get the current weather for a city.",
#     "parameters": {
#       "type": "object",
#       "properties": {"city": {"type": "string"}, "units": {"type": "string"}},
#       "required": ["city"]
#     }
#   }
# ]

# Always pass this to chat() — NOT a provider-specific schema:
response = llm.chat("What's the weather?", tools=registry.to_schema())

to_openai_schema() / to_anthropic_schema() / to_gemini_schema() — Per-provider (testing/debugging only)

These methods still exist and return provider-specific formats. They are useful for inspecting what a provider will receive on the wire, or for writing unit tests. Do not pass their output to LLM.chat(tools=...) — pass to_schema() instead.

registry.to_openai_schema()    # {"type": "function", "function": {...}}
registry.to_anthropic_schema() # {"name": ..., "input_schema": ...}
registry.to_gemini_schema()    # {"name": ..., "parameters": ...}

Other properties

registry.names   # ["get_weather", "search_docs"] — in registration order
len(registry)    # 2
repr(registry)   # "ToolRegistry(tools=['get_weather', 'search_docs'])"

---

Full Tool-Calling Flow

The complete lifecycle of a tool call, from definition to execution. Both layers (LLM, GraphLLM) use the same neutral format — only the registry call changes between providers.

from pywrapai       import LLM, Message, Role, ToolResult
from pywrapai_graph import GraphLLM, tool, ToolRegistry

llm   = LLM(provider="anthropic", model="claude-haiku-4-5-20251001")
graph = GraphLLM(llm=llm)

# Step 1 — Define and register tools
@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"16°C, cloudy in {city}"

@tool
def search_docs(query: str, limit: int = 5) -> str:
    """Search internal documentation."""
    return f"Found {limit} results for: {query}"

registry = ToolRegistry()
registry.register(get_weather)
registry.register(search_docs)

print(registry.names)   # ["get_weather", "search_docs"]
print(len(registry))    # 2

# Step 2 — Call the LLM with the NEUTRAL schema (to_schema(), not provider-specific)
response = graph.chat(
    "What's the weather in London?",
    tools       = registry.to_schema(),   # ← always use this
    tool_choice = "auto",
)

# Step 3 — Check whether the model wants to call a tool
if response.tool_calls:
    tc     = response.tool_calls[0]
    result = registry.execute(tc)
    print(result.content)   # "16°C, cloudy in London"

    # Step 4 — Build follow-up messages: assistant turn + tool result turn
    messages = [
        Message(Role.USER,      "What's the weather in London?"),
        Message(Role.ASSISTANT, "", tool_calls=response.tool_calls),
        Message(Role.TOOL,      result.content,
                tool_result=ToolResult(tc.id, tc.name, result.content)),
    ]

    # Step 5 — Call again so the model turns the result into a final reply
    final = graph.chat(messages)
    print(final.content)
    # → "The weather in London is 16°C and cloudy."

# The Agent in Phase 3 automates Steps 3–5 in a ReAct loop.

Provider-specific conversion is internal. When graph.chat(tools=registry.to_schema()) is called:

• OpenAI receives [{"type": "function", "function": {...}}]
• Anthropic receives [{"name": ..., "input_schema": {...}}]
• Gemini builds FunctionDeclaration objects with forbidden schema keys stripped
• Ollama follows OpenAI’s wire format; tool_choice is ignored with a warning

Your application code never changes between providers — only the LLM(provider=...) constructor.

---

History Injection in Detail

The history parameter of chat() is a list[Message] in chronological order. It is always injected between the system prompt and the current message.

Example — 3-turn conversation:

Turn 1 (no history):

[SYSTEM]    You are a helpful support agent.
[USER]      My name is Alice.
→ [ASSISTANT] Hello Alice! How can I help you?

Turn 2 (history has turn 1):

[SYSTEM]    You are a helpful support agent.
[USER]      My name is Alice.      ← history turn 1a
[ASSISTANT] Hello Alice! ...       ← history turn 1b
[USER]      What's my name?        ← current
→ [ASSISTANT] Your name is Alice.

Turn 3 (history has turns 1 and 2):

[SYSTEM]    You are a helpful support agent.
[USER]      My name is Alice.      ← history
[ASSISTANT] Hello Alice! ...       ← history
[USER]      What's my name?        ← history
[ASSISTANT] Your name is Alice.    ← history
[USER]      How many messages?     ← current
→ [ASSISTANT] You've sent 3 messages.

This is the only structure that gives the LLM the full context it needs to answer coherently.

---

Phase 2 — RAG Pipeline

The RAG pipeline is built from five composable primitives: Document, Chunker, InMemoryVectorStore, VectorRetriever, and GraphLLM(retriever=...). Each can be replaced independently.

Data Types

Document

from pywrapai_graph import Document

doc = Document(
    text     = "Paris is the capital of France and a major cultural hub.",
    metadata = {"source": "geography.txt", "page": 1},  # optional
)

doc.id        # str — auto-generated UUID
doc.text      # str
doc.metadata  # dict

Document represents a raw piece of text before chunking. You create Document objects from your own data sources (files, databases, web pages, etc.) and pass them to retriever.add_documents().

Chunk

from pywrapai_graph import Chunk

chunk = Chunk(
    text      = "Paris is the capital of France.",
    doc_id    = doc.id,
    metadata  = {"source": "geography.txt"},  # optional
)

chunk.id        # str — auto-generated UUID
chunk.text      # str
chunk.doc_id    # str — parent Document.id
chunk.metadata  # dict
chunk.embedding # list[float] | None — None until embedded

Chunk is the unit that gets embedded and stored. You do not normally create chunks manually — Chunker.split() and retriever.add_documents() handle this.

ScoredChunk

from pywrapai_graph import ScoredChunk

# Returned by VectorRetriever.retrieve() and InMemoryVectorStore.search()
scored = ScoredChunk(chunk=chunk, score=0.87)

scored.chunk  # Chunk
scored.score  # float — cosine similarity, 0.0–1.0 (higher = more similar)

Results from retrieve() are sorted by score descending. The top-k are injected into the prompt.

---

Chunker

from pywrapai_graph import Chunker

chunker = Chunker(
    chunk_size = 500,   # target character count per chunk
    overlap    = 50,    # characters of overlap between adjacent chunks
)

chunks: list[Chunk] = chunker.split(doc)

Splitting priority order:

1. Paragraph boundaries (\n\n) — splits at blank lines first
2. Sentence boundaries (., !, ? followed by whitespace) — if a paragraph is too large
3. Hard character cut — last resort if no sentence boundary found within chunk_size

overlap copies the last n characters of each chunk to the beginning of the next, preserving cross-boundary context. Chunks shorter than 10 characters are dropped.

You rarely call Chunker.split() directly — retriever.add_documents(docs, chunker=chunker) handles chunking and embedding in one call.

---

InMemoryVectorStore

from pywrapai_graph import InMemoryVectorStore

store = InMemoryVectorStore()

An in-memory vector store using cosine similarity search (numpy L2-normalized dot product).

# Add pre-embedded chunks
store.add(chunks)   # all chunks must have chunk.embedding set

# Search
results: list[ScoredChunk] = store.search(query_vector, k=4)

# Inspect
len(store)      # number of stored chunks
store.clear()   # remove all chunks

add() raises ValueError if any chunk has embedding=None. Always embed chunks before adding.

search() returns at most k results sorted by cosine similarity (highest first). Returns an empty list if the store is empty.

Implement BaseVectorStore to add your own backends (Pinecone, Weaviate, Milvus, etc.):

from pywrapai_graph import BaseVectorStore, ScoredChunk, Chunk

class PineconeStore(BaseVectorStore):
    def add(self, chunks: list[Chunk]) -> None: ...
    def search(self, query_vector: list[float], k: int) -> list[ScoredChunk]: ...
    def clear(self) -> None: ...
    def __len__(self) -> int: ...

---

Persistent Vector Store Backends

All backends implement BaseVectorStore and are swappable without changing any other code. VectorRetriever and GraphLLM see only the BaseVectorStore interface.

Comparison table

Class	Install	Infra	Scale	Persistence
InMemoryVectorStore	numpy (always)	None	~100k chunks	No
SQLiteVectorStore	stdlib only	Single .db file	~100k chunks	Yes
ChromaVectorStore	pip install chromadb	None / Docker	Millions	Yes
FAISSVectorStore	pip install faiss-cpu	Two local files	Billions (exact: ~1M)	Yes
PGVectorStore	pip install psycopg2-binary pgvector	PostgreSQL server	Hundreds of millions	Yes
QdrantVectorStore	pip install qdrant-client	None / Docker / Cloud	Billions	Yes

---

SQLiteVectorStore

Zero-dependency persistent vector store. Stores chunk text and metadata in a SQLite table; embedding vectors as BLOB columns (serialised numpy float32). On search, all rows are loaded into a numpy matrix and searched with the same dot-product algorithm as InMemoryVectorStore. The difference from InMemoryVectorStore is purely persistence — vectors survive process restarts.

from pywrapai_graph import SQLiteVectorStore

store = SQLiteVectorStore(
    db_path    = "./knowledge_base.db",  # file is created if it doesn't exist
    table_name = "vector_chunks",        # default
)

---

ChromaVectorStore

Chromadb embedded vector database. HNSW approximate nearest-neighbor indexing. No server required in embedded mode. Scales to millions of vectors.

from pywrapai_graph import ChromaVectorStore

# Embedded (local directory)
store = ChromaVectorStore(path="./chroma_db", collection="my_docs")

# HTTP server (requires running chromadb server)
store = ChromaVectorStore(url="http://localhost:8000", collection="my_docs")

Exactly one of path or url must be provided.

---

FAISSVectorStore

Facebook FAISS flat inner-product index. Vectors are L2-normalised so inner product equals cosine similarity. FAISS stores only raw floats — chunk metadata (text, doc_id, metadata dict) lives in a parallel pickle file.

from pywrapai_graph import FAISSVectorStore

store = FAISSVectorStore(
    index_path = "./docs.faiss",     # FAISS index; created on first add()
    meta_path  = "./docs_meta.pkl",  # parallel chunk metadata; created on first add()
    dimension  = None,               # inferred from first add() — or pass explicitly
)
store.add(chunks)    # saves both files immediately
results = store.search(query_vec, k=4)

If both files exist at startup they are loaded automatically — no re-embedding on restart.

---

PGVectorStore

PostgreSQL + pgvector extension. Uses the <=> cosine-distance operator for server-side nearest-neighbour search. The target store for Stage 2 multi-server deployments.

from pywrapai_graph import PGVectorStore

store = PGVectorStore(
    connection_string = "postgresql://user:pass@localhost:5432/mydb",
    table_name        = "vector_chunks",  # default
    dimension         = 1536,             # must match your embedding model
)
store.add(chunks)
results = store.search(query_embedding, k=4)
store.close()   # closes the psycopg2 connection

Requirement: PostgreSQL must have the pgvector extension available. The table and extension are created automatically on first instantiation.

Live tests: Set PYWRAPAI_TEST_PG_URL=postgresql://user:pass@host/db to run PGVectorStore tests against a real server; otherwise they are skipped automatically.

---

QdrantVectorStore

Qdrant dedicated vector database. Supports in-process memory, local disk, Docker self-hosting, and Qdrant Cloud all with the same constructor API.

from pywrapai_graph import QdrantVectorStore

# In-process memory (tests, experimentation)
store = QdrantVectorStore(collection="my_docs", location=":memory:", dimension=1536)

# Local disk — survives restarts, no server required
store = QdrantVectorStore(collection="my_docs", location="./qdrant_data")

# Docker / self-hosted
store = QdrantVectorStore(collection="my_docs", url="http://localhost:6333")

# Qdrant Cloud
store = QdrantVectorStore(collection="my_docs", url="https://xyz.qdrant.io", api_key="...")

Important — local disk mode holds an exclusive file lock. Call store.close() before opening a second instance pointing to the same path:

store1 = QdrantVectorStore("docs", location="./qdrant_data")
store1.add(chunks)
store1.close()   # release the lock

store2 = QdrantVectorStore("docs", location="./qdrant_data")
results = store2.search(query_vec)

---

VectorRetriever

VectorRetriever connects an Embedder to a BaseVectorStore. It embeds queries, searches the store, and optionally filters by minimum similarity score.

from pywrapai       import Embedder
from pywrapai_graph import InMemoryVectorStore, VectorRetriever

embedder = Embedder(provider="openai", model="text-embedding-3-small")
store    = InMemoryVectorStore()

retriever = VectorRetriever(
    embedder  = embedder,
    store     = store,
    k         = 4,          # default chunks to retrieve
    min_score = 0.0,        # minimum cosine similarity (0.0 = no filter)
)

add_documents(docs, chunker=None)

from pywrapai_graph import Document, Chunker

docs = [
    Document("Paris is the capital of France."),
    Document("The Eiffel Tower is 330 metres tall."),
]

chunker = Chunker(chunk_size=500, overlap=50)
n_chunks = retriever.add_documents(docs, chunker=chunker)
# n_chunks → 2 (one chunk per short document in this example)

What happens inside:

1. Each Document is split into Chunk objects using chunker.split() (or no splitting if chunker=None)
2. All chunks are embedded in one batched Embedder.embed() call
3. Chunks with their embedding set are added to the store

Returns the total number of chunks added.

If chunker=None, each Document becomes exactly one Chunk (no splitting). Use this when your documents are already pre-chunked.

retrieve(query, k=None)

results: list[ScoredChunk] = retriever.retrieve("What is the capital of France?", k=3)

for sc in results:
    print(f"[{sc.score:.3f}] {sc.chunk.text}")
# [0.91] Paris is the capital of France.
# [0.43] The Eiffel Tower is 330 metres tall.

Embeds query with one Embedder.embed_one() call, searches the store, applies min_score filter, and returns up to k results (defaults to the k set at construction).

Implement BaseRetriever to add alternative retrieval strategies (BM25, hybrid search, graph memory, etc.):

from pywrapai_graph import BaseRetriever, ScoredChunk

class BM25Retriever(BaseRetriever):
    def retrieve(self, query: str, k: int = 4) -> list[ScoredChunk]: ...
    def add_documents(self, docs, chunker=None) -> int: ...

---

retrieval_tool()

from pywrapai_graph import retrieval_tool, ToolRegistry

registry = ToolRegistry()

# Wrap a retriever as a callable Tool for the agent loop (Phase 3)
t = retrieval_tool(
    retriever   = retriever,
    name        = "search_knowledge_base",
    description = "Search the knowledge base for relevant facts.",
)

registry.register(t)

# The LLM can now call "search_knowledge_base" as a tool
response = graph.chat("Find information about the Eiffel Tower",
                      tools=registry.to_schema())

retrieval_tool() returns a Tool dataclass whose fn calls retriever.retrieve(query) and formats the results as a numbered list string:

[1] Paris is the capital of France.
[2] The Eiffel Tower is 330 metres tall and located in Paris.

This is the bridge between the RAG pipeline and the Phase 3 agent loop — it makes any BaseRetriever directly callable by the LLM as a tool.

---

Full RAG Example

from pywrapai       import LLM, Embedder
from pywrapai_graph import (
    GraphLLM, Document, Chunker,
    InMemoryVectorStore, VectorRetriever,
)

# 1. Build the knowledge base
embedder  = Embedder(provider="openai", model="text-embedding-3-small")
store     = InMemoryVectorStore()
retriever = VectorRetriever(embedder=embedder, store=store, k=4)

docs = [
    Document("Paris is the capital of France, founded in 987 AD."),
    Document("The Louvre museum houses over 380,000 works of art."),
    Document("France uses the Euro as its currency since 2002."),
]
retriever.add_documents(docs, chunker=Chunker(chunk_size=500))

# 2. Build the graph with RAG
llm   = LLM(provider="openai")
graph = GraphLLM(llm=llm, retriever=retriever, rag_k=3)

# 3. Chat — RAG context is injected automatically
response = graph.chat("What is the capital of France?")
print(response.content)   # "The capital of France is Paris, founded in 987 AD."

# 4. Skip RAG for a specific call
response = graph.chat("Thank you!", rag=False)  # no retrieval, plain chat

# 5. Check embedding cost separately
print(embedder.tokens.total_cost)  # USD cost of all embed() calls
print(llm.tokens.total_cost)       # USD cost of all chat() calls

D4 Template (RAG context injection format)

When chunks are retrieved, the current user message is replaced with this template:

Use the following context to answer the question.

<context>
[1] {chunk_1_text}
[2] {chunk_2_text}
[3] {chunk_3_text}
</context>

Question: {original_user_message}

Design decisions:

• D4: RAG context is injected into the text of the current user message — not as a separate message. This keeps message count constant and prevents the context from leaking into history.
• D5: Only the original user message should be saved to history. The <context> block is stripped before persistence. This means follow-up queries always retrieve fresh context rather than relying on stale chunks from prior turns.
• D6: Gemini and Ollama embed APIs do not return token counts. A 4-characters-per-token heuristic is used for these providers. OpenAI returns exact counts.
• D7: Embedding costs are tracked as CallRecord(kind="embedding") in the Embedder’s own TokenTracker. They are separate from the LLM’s TokenTracker but can be read at any time via embedder.tokens.

---

System Prompt Versioning

The versioning model is simple: the version is part of the filename.

support-agent-v1.txt and support-agent-v2.txt are two independent prompts. You decide which version to use at session creation time, and it stays fixed for the lifetime of that session.

# Version A — using v1
graph_v1 = GraphLLM(
    llm         = llm,
    enable_system_prompt_versioning = True,
    prompt_file = "support-agent-v1.txt",
)

# Version B — using v2 for A/B testing
graph_v2 = GraphLLM(
    llm         = llm,
    enable_system_prompt_versioning = True,
    prompt_file = "support-agent-v2.txt",
)

Since a persistence layer built on top of GraphLLM can record session_id alongside every response, you can later compare which prompt version produced better outcomes (lower cost, fewer turns to resolution, higher satisfaction scores, etc.).

---

Practical Patterns

Pattern 1 — Stateless (manual history)

Suitable for applications that manage their own session state or that call GraphLLM from a stateless request handler.

# Application fetches history from its own DB
history = fetch_history_from_my_db(user_id, session_id)

response = graph.chat("Where is my order?", history=history)

# Application saves the response to its own DB
save_to_my_db(user_id, session_id, "user",      "Where is my order?")
save_to_my_db(user_id, session_id, "assistant", response.content)

Pattern 2 — No system prompt versioning

When you want to set the system prompt inline rather than from a file:

llm = LLM(
    provider      = "anthropic",
    system_prompt = "You are a concise code reviewer. Review code for bugs and style issues.",
)

graph = GraphLLM(llm=llm)   # enable_system_prompt_versioning=False (default)

response = graph.chat("def add(a, b): return a + c")

---

Design Decisions

Why a separate library? PyWrapAI is the engine — it calls the API. PyWrapAI-Graph is the orchestrator — it decides what to send. Separating them means you can use PyWrapAI alone for simple one-shot calls (scripts, data pipelines, batch jobs) without pulling in the message assembly machinery. And it means PyWrapAI-Graph can evolve independently — adding RAG, agents, and workflow nodes — without touching the API-calling core.

Why does graph.chat() not store history automatically? That’s the job of whatever persistence layer you put on top of it. PyWrapAI-Graph accepts history as a parameter; it does not own the database. This separation means you can use PyWrapAI-Graph without any database at all (useful in tests, scripts, or when you handle persistence yourself).

Why is the system prompt always position 0? All major LLM APIs expect the system prompt at the beginning of the message array. Putting it anywhere else produces undefined behaviour. PyWrapAI-Graph enforces this by design and silently removes any system message from the caller’s input list when versioning is enabled, preventing duplicate system prompts.

Why does @tool live in pywrapai_graph, not pywrapai? pywrapai is the engine — it fires API calls and returns responses. It has no concept of a function registry or execution layer. Tool registration and execution are orchestration concerns that belong in pywrapai_graph.

Why does @tool use _pywrapai_tool instead of marking the function itself? The decorator must preserve the original function’s identity so it still works as a normal callable. Storing the Tool metadata in a separate attribute (wrapper._pywrapai_tool) keeps the two concerns cleanly separated: the callable and its schema metadata. ToolRegistry.register() checks for this attribute to distinguish decorated functions from plain ones.

Why does execute() catch exceptions instead of raising? When a tool crashes at runtime (network timeout, invalid argument, external API failure), crashing the entire agent loop discards context and forces the user to restart. Returning the exception message as a string lets the model see what went wrong and attempt recovery — retry, ask for clarification, or call a different tool. Production agents need graceful degradation.

Why to_schema() instead of calling a per-provider method? Phase 0 and Phase 1 (items 1.1–1.5) added three provider-specific schema methods (to_openai_schema(), to_anthropic_schema(), to_gemini_schema()). Phase 1.6 added to_schema() as the single neutral format that callers always use. Each provider now converts internally — the format difference (OpenAI wraps in {"type": "function", "function": {...}}, Anthropic uses input_schema, Gemini builds FunctionDeclaration objects) is an internal implementation detail, not something the caller should need to know about. One format, four providers. The per-provider methods still exist for inspection and unit testing, but application code never calls them directly.

Why does stream(tools=...) raise ValueError? Streaming with tool calls produces interleaved delta events — partial function names, partial JSON argument fragments — that must be assembled across chunks before execution. That assembly logic belongs in the Phase 3 agent loop, not in the streaming path. Making it ValueError at both layers (LLM, GraphLLM) gives a clear error message rather than silently dropping tool calls or producing corrupted output.

Why does with_system_prompt() use a shallow copy rather than a setter? A typical deployment creates one GraphLLM at startup and reuses it across all users and sessions, and each user may need a different system prompt. A setter would mutate the shared instance and cause race conditions under concurrent requests. A shallow copy returns a separate object per user while deliberately sharing the _llm instance — so the cache and TokenTracker remain unified. A deep copy would sever that sharing, duplicating cache entries and splitting cost accounting per user, which is never what you want.

Why does Embedder live in pywrapai, not pywrapai_graph? (D1) Embedding is a provider API call — it calls OpenAI, Gemini, or Ollama over HTTP and returns a vector. That is the same kind of work as LLM.chat(). The pywrapai layer owns all provider API calls. The RAG pipeline (chunking, storing, retrieving, injecting into messages) lives in pywrapai_graph because that is orchestration.

Why is RAG context injected into the message text rather than added as a separate message? (D4) Adding a separate context message would increase the message count and appear in history if the application stores all messages. Injecting into the user message text keeps the message structure identical to non-RAG calls. History still stores just the original user query — the <context> block is never persisted (D5).

Why is RAG context never saved to history? (D5) Follow-up queries should retrieve fresh context relevant to the new question, not inherit stale chunks from the previous turn. Persisting RAG context would also inflate the stored message size significantly — a 4-chunk context block could be 2,000+ characters per turn.

Why is retrieval_tool() a free function rather than a method on BaseRetriever? BaseRetriever is in pywrapai_graph.rag.retriever. Tool is in pywrapai_graph.tools.tool. Making BaseRetriever depend on Tool would create an intra-library circular import. The free function retrieval_tool() imports Tool at call time and bridges the two submodules cleanly.

---

Phase 3 — Agent ReAct Loop

Agent drives the full Reason-Act (ReAct) loop automatically: it calls the LLM with tool definitions, executes the tool calls the model requests, appends the results, and repeats until the model returns a plain-text answer or a turn limit is reached.

File: pywrapai_graph/agent.py

---

Types

AgentConfig

Controls loop behaviour. Pass to Agent() at construction time.

from pywrapai_graph import AgentConfig

config = AgentConfig(
    max_turns         = 10,               # stop after this many tool-calling turns
    tool_error_policy = "return_to_model", # what to do when a tool raises an exception
    on_step           = None,             # optional callback fired after each completed turn
)

tool_error_policy values:

Value	Behaviour
"return_to_model" (default)	Catches the exception, formats it as ToolResult(content="Error: <str(exc)>"), and continues. The model sees what went wrong and can recover.
"raise"	Lets the exception propagate immediately, aborting the run.
"skip"	Silently drops the failed tool’s result. The model’s next turn does not receive a TOOL message for that call. Use with care.

KeyError (unregistered tool name) always propagates regardless of policy — it is a programmer error, not a runtime failure.

on_step callback:

def log_step(step: AgentStep) -> None:
    print(f"Turn {step.turn_num}: {len(step.tool_calls)} call(s), cost=${step.cost_usd:.6f}")

config = AgentConfig(on_step=log_step)

The callback fires synchronously after _execute_tools() returns for that turn, before the step is appended to the internal list. All fields of AgentStep are fully populated at callback time. If the callback raises, the exception propagates out of run().

---

AgentStep

One completed turn of the loop. Returned in AgentResult.steps and received by on_step.

step.turn_num      # int   — 1-indexed turn number
step.tool_calls    # list[ToolCall]   — what the model requested
step.tool_results  # list[ToolResult] — what your functions returned
step.usage         # TokenUsage — input/output tokens for this turn's LLM call
step.cost_usd      # float — USD cost for this turn

---

AgentResult

Returned by Agent.run() and Agent.arun() on success.

result.content         # str         — the model's final plain-text answer
result.steps           # list[AgentStep] — one entry per tool-calling turn
result.total_turns     # int         — total LLM calls made (including the final answer turn)
result.total_usage     # TokenUsage  — sum of all turns' input/output tokens
result.total_cost_usd  # float       — sum of all turns' cost_usd

The final plain-text answer turn (no tool calls) is included in total_turns and total_usage, but does not produce an AgentStep (steps contain only tool-calling turns).

# Invariant:
assert result.total_cost_usd == sum(s.cost_usd for s in result.steps) + final_turn_cost

---

AgentMaxTurnsError

Raised when the loop reaches max_turns without getting a plain-text answer.

from pywrapai_graph import AgentMaxTurnsError

try:
    result = agent.run(input)
except AgentMaxTurnsError as exc:
    print(f"Loop ran {exc.turns} turns without finishing.")
    partial = exc.partial   # AgentResult with whatever completed
    print(partial.total_cost_usd)   # cost up to the limit

exc.partial is a complete AgentResult with all completed steps and accumulated cost. Use it to inspect what the agent did before hitting the limit, or to return a partial answer to the user.

---

Agent — Constructor

from pywrapai       import LLM
from pywrapai_graph import GraphLLM, ToolRegistry, tool, Agent, AgentConfig

llm      = LLM(provider="openai", model="gpt-4o-mini")
graph    = GraphLLM(llm=llm).with_system_prompt("You are a helpful assistant.")
registry = ToolRegistry()

@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"16°C, cloudy in {city}"

registry.register(get_weather)

config = AgentConfig(max_turns=8, tool_error_policy="return_to_model")
agent  = Agent(graph=graph, registry=registry, config=config)

Parameters:

Parameter	Type	Description
graph	GraphLLM	The assembled LLM + system prompt. Create with .with_system_prompt()
registry	ToolRegistry	All tools the model is allowed to call
config	AgentConfig	Loop configuration — max turns, error policy, step callback

---

Agent.run() — Synchronous

result: AgentResult = agent.run(input)

Parameters:

• input: str | list[Message] — the user’s question or a pre-built message list
• history: list[Message] | None — optional prior conversation context injected into every GraphLLM.chat() call

Loop behaviour:

messages = [current input]
for turn in range(1, max_turns + 1):
    response = graph.chat(messages, tools=registry.to_schema(), rag=False)
    record usage → AgentStep

    if response.tool_calls is None or empty:
        return AgentResult(content=response.content, ...)

    results  = _execute_tools(response.tool_calls, policy)
    messages += [ASSISTANT + tool_calls] + [TOOL per result]
    steps.append(AgentStep(...))
    if on_step: on_step(steps[-1])

raise AgentMaxTurnsError(turn, partial_result)

rag=False is always passed — retrieval and tool-calling are separate concerns. The message list grows with each turn so the model retains the full reasoning history. Token usage from response.usage is accumulated into AgentResult.total_usage.

---

Agent.arun() — Async

result: AgentResult = await agent.arun(input)

Same behaviour as run(), using graph.achat() (the provider’s native async client). When the model returns multiple tool calls in one response, they are dispatched concurrently with asyncio.gather — all tools in a turn run at the same time.

import asyncio

async def main():
    result = await agent.arun("What is the weather in London and Paris?")
    print(result.content)
    print(f"Total cost: ${result.total_cost_usd:.6f}")

asyncio.run(main())

Parallel execution note: Tool functions are currently synchronous (registry.execute() is sync). asyncio.gather dispatches them in the event loop but they run sequentially within it — they do not yield. True async tool support (tools that call external APIs with await) is a future enhancement.

---

Full Agent Example

from pywrapai       import LLM
from pywrapai_graph import (
    GraphLLM, ToolRegistry, tool,
    Agent, AgentConfig, AgentResult, AgentMaxTurnsError,
)

# 1. Define tools
@tool
def search_docs(query: str) -> str:
    """Search internal documentation for a query."""
    return f"Found 3 results for '{query}': ..."

@tool
def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"16°C, cloudy in {city}"

# 2. Register
registry = ToolRegistry()
registry.register(search_docs)
registry.register(get_weather)

# 3. Build agent
llm    = LLM(provider="openai", model="gpt-4o-mini")
graph  = GraphLLM(llm=llm).with_system_prompt(
    "You are a helpful assistant. Always call a tool to look up real data."
)

def log_step(step):
    for tc in step.tool_calls:
        print(f"  [{step.turn_num}] called {tc.name}({tc.arguments})")

config = AgentConfig(max_turns=8, on_step=log_step)
agent  = Agent(graph=graph, registry=registry, config=config)

# 4. Run
try:
    result: AgentResult = agent.run("What's the weather in London?")
    print(result.content)
    print(f"Turns: {result.total_turns}  Cost: ${result.total_cost_usd:.6f}")

except AgentMaxTurnsError as exc:
    print(f"Hit limit after {exc.turns} turns.")
    print(f"Partial cost: ${exc.partial.total_cost_usd:.6f}")

---

Agent Design Decisions

Why rag=False inside Agent.run()? RAG retrieval is automatic and query-based. In a tool-calling loop the user’s original question may not be the right retrieval query by turn 3 — the model has already called tools and received information. Mixing automatic RAG injection with tool-calling creates ambiguous prompts and unpredictable context. If you want the agent to retrieve documents, give it a retrieval_tool() — it decides when to retrieve, which query to use, and how to combine retrieved chunks with other tool results.

Why does _execute_tools call registry.get(name).fn() directly instead of registry.execute()? ToolRegistry.execute() already silently catches all exceptions and returns them as ToolResult.content. That is the right behaviour for single-shot tool use. But Agent owns the error-handling layer — it must be able to apply tool_error_policy="raise" or "skip" by catching the exception itself. Calling execute() would swallow all exceptions before Agent could see them.

Why does the message accumulator grow with each turn? The model needs the full tool call / result history to reason correctly. Without it, the model can’t see what it already tried or what the previous tool returned — each turn would look like a fresh start. This means token usage grows quadratically in long agent runs; max_turns exists to cap that growth.

Why is AgentMaxTurnsError raised rather than returning a partial result? Returning a partial AgentResult would silently give the caller an incomplete answer with no indication that it is incomplete. A raised exception forces the caller to handle the partial case explicitly — either by surfacing the error to the user or by reading exc.partial and deciding what to do. Silence is wrong here.