PyWrapAI — Architecture

This document explains why the ecosystem is split into separate libraries and how they interact at runtime.

---

The Two-Library Design

Library	Replaces	Single Responsibility
PyWrapAI	LangChain	Call an LLM and return a response
PyWrapAI-Graph	LangGraph	Assemble messages and manage prompts

These are separately installable Python packages with a strict, one-way dependency graph:

PyWrapAI-Graph  →  imports  →  PyWrapAI

• PyWrapAI imports nothing from PyWrapAI-Graph
• PyWrapAI-Graph imports from PyWrapAI only

This directionality is enforced by design. Adding a dependency in the other direction (PyWrapAI importing from PyWrapAI-Graph) would collapse the separation and force users of the simple LLM-calling library to install the full stack.

---

What Each Library Owns

PyWrapAI — The Engine

Knows about: LLM providers, API calls, tokens, cost, caching, retries, tool wire format, embeddings.

Does not know about: conversation structure, history, system prompts, databases, analytics, tool registration, RAG pipelines.

LLM(provider, model, ...)
  ├── chat(messages, tools=, tool_choice=)   → LLMResponse
  ├── stream(messages)                       → Iterator[str]  (tools= raises ValueError)
  ├── chat_structured(...)                   → Pydantic model | dict
  ├── achat(messages, tools=, tool_choice=)  → awaitable LLMResponse
  ├── astream(messages)                      → AsyncIterator[str]
  └── tokens                                 → TokenTracker

Embedder(provider, model, track_tokens=True)
  ├── embed(texts)       → list[list[float]]   (batched; one API call per 100 texts)
  ├── embed_one(text)    → list[float]
  ├── aembed(texts)      → awaitable list[list[float]]
  └── tokens             → TokenTracker        (kind="embedding" CallRecords)

PyWrapAI-Graph — The Assembler

Knows about: message assembly order, system prompt files, history injection, tool registration and execution, RAG retrieval and context injection, ReAct agent loop.

Does not know about: databases, user identities, session IDs, analytics.

GraphLLM(llm, prompt_file=, retriever=None, rag_k=4, ...)
  ├── chat(input, history, tools=, tool_choice=, rag=True)  → LLMResponse
  ├── stream(input, history, rag=True)                      → Iterator[str]  (tools= raises ValueError)
  ├── chat_structured(input, schema, history, rag=True)     → model | dict
  └── achat(input, history, tools=, tool_choice=, rag=True) → awaitable LLMResponse

ToolRegistry
  ├── register(@tool-decorated fn or Tool)
  ├── to_schema()              → list[dict]  ← neutral format, always use this
  ├── to_openai_schema()       → list[dict]  ← for inspection/testing only
  ├── to_anthropic_schema()    → list[dict]  ← for inspection/testing only
  ├── to_gemini_schema()       → list[dict]  ← for inspection/testing only
  └── execute(ToolCall)        → ToolResult

Agent(graph, registry, config)          ← Phase 3 ReAct loop
  ├── run(input, history=None)          → AgentResult   (sync)
  └── arun(input, history=None)         → awaitable AgentResult

AgentConfig(max_turns=10, tool_error_policy="return_to_model", on_step=None)
AgentStep(turn_num, tool_calls, tool_results, usage, cost_usd)
AgentResult(content, steps, total_turns, total_usage, total_cost_usd)
AgentMaxTurnsError(turns, partial: AgentResult)

Phase 2 — RAG Pipeline

Document(text, metadata={})           ← text + optional metadata dict; auto-generates id
Chunk(text, doc_id, metadata={})      ← chunk of a Document; embedding=None until embedded
ScoredChunk(chunk, score)             ← chunk + cosine similarity score from search

Chunker(chunk_size=500, overlap=50)
  └── split(doc) → list[Chunk]        ← paragraph > sentence > hard-cut boundary priority

BaseVectorStore (ABC)                 ← implement to add Pinecone, Weaviate, etc.
  ├── InMemoryVectorStore()           ← numpy L2-norm dot product; RAM only; no deps beyond numpy
  ├── SQLiteVectorStore(db_path)      ← stdlib sqlite3; persisted; same scale ceiling as in-memory
  ├── ChromaVectorStore(path, collection)  ← HNSW ANN; embedded or HTTP server; pip install chromadb
  ├── FAISSVectorStore(index_path, meta_path)  ← exact IP index; parallel .pkl metadata; pip install faiss-cpu
  ├── PGVectorStore(connection_string)     ← PostgreSQL + pgvector; <=> cosine distance; pip install psycopg2-binary pgvector
  └── QdrantVectorStore(collection, location=|url=)  ← HNSW; :memory:/local/Docker/Cloud; pip install qdrant-client

BaseRetriever (ABC)                   ← implement to add BM25, graph memory, etc.
VectorRetriever(embedder, store, k=4, min_score=0.0)
  ├── retrieve(query, k=None)        → list[ScoredChunk]
  └── add_documents(docs, chunker=None) → int  (chunk → embed → store; one batched embed call)

retrieval_tool(retriever, name="retrieve", description="...") → Tool
  └── wraps any BaseRetriever as a Tool for use in agent loops (Phase 3)

---

Runtime Message Flow

Plain text chat

When a user sends a message, here is the exact sequence of operations:

User types: "Where is my order?"

Application loads history from wherever it persists it
  → [USER: "I ordered a laptop", ASSISTANT: "Order #12345 found"]

GraphLLM.chat("Where is my order?", history=[...])
  │
  ├─ 1. Assemble full message list:
  │       [SYSTEM] You are a helpful support agent.
  │       [USER]   I ordered a laptop.              ← history
  │       [ASST]   Order #12345 found.             ← history
  │       [USER]   Where is my order?              ← current
  │
  └─ 2. LLM.chat(assembled_messages)
          │
          ├─ 3. Check cache → miss
          ├─ 4. Call provider API
          ├─ 5. Record token usage in TokenTracker
          └─ 6. Return LLMResponse (content="Your order is in transit")

Application appends the new turn to its own history store and
returns the LLMResponse

RAG chat (Phase 2)

When a retriever= is set on GraphLLM, the current user message is automatically expanded with retrieved context before the LLM call:

GraphLLM.chat("What is the capital of France?", history=[...], rag=True)  # retriever= set on GraphLLM
  │
  ├─ 1. Retrieve from VectorRetriever:
  │       embedder.embed_one("What is the capital of France?")  → query vector
  │       store.search(query_vec, k=4)  → list[ScoredChunk]
  │
  ├─ 2. Inject RAG context into current message (D4 template):
  │       "Use the following context to answer the question.
  │
  │        <context>
  │        [1] Paris is the capital of France and a major European city.
  │        </context>
  │
  │        Question: What is the capital of France?"
  │
  ├─ 3. Assemble full message list:
  │       [SYSTEM] You are a helpful assistant.
  │       [USER]   ... (history) ...
  │       [ASST]   ... (history) ...
  │       [USER]   Use the following context...  ← RAG-augmented current message
  │
  └─ 4. LLM.chat(assembled_messages)

If the application persists history, it should save only the original
user message "What is the capital of France?" — the RAG context is
never saved (design decision D5), so follow-up queries always retrieve
fresh context.

Tool-calling chat (Phase 1.6)

When tools are passed, the LLM may return a tool call instead of text. The agent loop (Phase 3) automates this, but here is the explicit flow:

GraphLLM.chat("What's the weather?", tools=registry.to_schema())
  │
  ├─ (Steps 1–5 same as above)
  │
  ├─ 6. LLM.chat() converts neutral tools to provider wire format internally:
  │       OpenAI   → [{"type": "function", "function": {...}}]
  │       Anthropic → [{"name": ..., "input_schema": {...}}]
  │       Gemini   → FunctionDeclaration objects (forbidden keys stripped)
  │       Ollama   → OpenAI format (tool_choice ignored with warning)
  │
  ├─ 7. Return LLMResponse(content="", tool_calls=[ToolCall(id=..., name="get_weather", ...)])
  │
  ├─ 8. Application runs:  result = registry.execute(response.tool_calls[0])
  │
  ├─ 9. Application builds follow-up messages:
  │       [ASST]  "" + tool_calls=[...]    ← assistant's decision to call a tool
  │       [TOOL]  "16°C, cloudy"           ← your function's result
  │
  └─ 10. Call chat() again with the follow-up messages
        → LLMResponse(content="The weather in London is 16°C and cloudy.")

From the application’s perspective, that final response is all that matters. The Agent class (Phase 3) automates this entire loop.

Agent ReAct loop (Phase 3)

When Agent.run() is called, the loop runs entirely automatically:

agent.run("What's the weather in London?")
  │
  ├─ turn 1: graph.chat(messages, tools=registry.to_schema(), rag=False)
  │     → LLMResponse(tool_calls=[ToolCall(name="get_weather", arguments={"city":"London"})])
  │
  ├─ _execute_tools([ToolCall(...)], policy="return_to_model")
  │     → [ToolResult(content="16°C, cloudy in London")]
  │
  ├─ _build_tool_messages(response, results)
  │     → [Message(ASSISTANT, "", tool_calls=[...]), Message(TOOL, "16°C, cloudy...")]
  │
  ├─ messages += [ASSISTANT] + [TOOL]
  │   AgentStep(turn_num=1, tool_calls=[...], tool_results=[...], usage=..., cost=...)
  │   on_step(steps[-1])   ← application callback fires here
  │
  ├─ turn 2: graph.chat(messages, tools=registry.to_schema(), rag=False)
  │     → LLMResponse(content="The weather in London is 16°C and cloudy.", tool_calls=None)
  │
  └─ return AgentResult(
         content       = "The weather in London is 16°C and cloudy.",
         steps         = [AgentStep(turn_num=1, ...)],
         total_turns   = 2,        ← 1 tool turn + 1 answer turn
         total_usage   = TokenUsage(input=..., output=...),
         total_cost_usd = 0.000123
     )

---

The Full Message Structure

In a production AI application, every LLM call contains exactly this structure:

┌─────────────────────────────────────────────┐
│  SYSTEM PROMPT                               │  ← position 0, always
│  (loaded from system_prompts_registry/)      │
├─────────────────────────────────────────────┤
│  HISTORY TURN 1 — USER                      │  ← oldest message
│  HISTORY TURN 1 — ASSISTANT                 │
│  HISTORY TURN 2 — USER                      │
│  HISTORY TURN 2 — ASSISTANT                 │
│  ...                                         │
│  HISTORY TURN N — USER                      │
│  HISTORY TURN N — ASSISTANT                 │  ← most recent stored turn
├─────────────────────────────────────────────┤
│  CURRENT USER MESSAGE                        │  ← what the user just typed
└─────────────────────────────────────────────┘

Why this order matters:

• System prompt at position 0 is required by all major LLM APIs
• History in chronological order allows the model to track conversation flow
• Current message at the end is what the model is being asked to respond to
• RAG context is injected into the current message text (not as a separate message), keeping history clean

RAG-augmented message structure:

┌─────────────────────────────────────────────┐
│  SYSTEM PROMPT                               │  ← position 0, always
├─────────────────────────────────────────────┤
│  HISTORY TURN 1 — USER                      │
│  HISTORY TURN 1 — ASSISTANT                 │
│  ...                                         │
├─────────────────────────────────────────────┤
│  CURRENT USER MESSAGE (RAG-augmented)        │  ← D4 template wraps the original query
│  "Use the following context...               │
│   <context>                                  │
│   [1] chunk text                             │
│   </context>                                 │
│   Question: {original query}"                │
└─────────────────────────────────────────────┘

Only the original query ("What is the capital of France?") should be persisted — the RAG context is NOT saved to history (design decision D5).

---

What Is NOT in History

These are explicit design decisions, not limitations:

Not stored	Why
System prompt	Sent fresh at position 0 every call. Storing it would duplicate tokens
RAG context	Retrieved fresh for each query via Embedder + VectorStore. Persisting it would waste tokens on follow-up turns where the same chunks may be irrelevant. The original user query IS stored — only the injected <context> block is dropped (design decision D5)
Tool call results (in reconstructed history)	Whatever persistence layer sits above GraphLLM should reload only plain user/assistant turns as Message objects. Each new agent run gets a clean dialogue history, not a dump of every JSON tool payload

---

Separation of Concerns — A Worked Example

Consider a developer who wants to:

1. Add response caching → change LLM(cache=True) — touches only PyWrapAI
2. Switch from v1 to v2 system prompt → change prompt_file= — touches only PyWrapAI-Graph
3. Switch from Anthropic to OpenAI → change provider= — touches only PyWrapAI
4. Add RAG over a new document set → wire up a retriever= — touches only PyWrapAI-Graph

No change ever touches more than one library. This is the design working as intended.

---

Provider Architecture

Every LLM provider implements the same BaseLLM interface:

BaseLLM (abstract)
  ├── AnthropicProvider
  ├── OpenAIProvider
  ├── GeminiProvider
  └── OllamaProvider

All providers are registered in pywrapai/providers/__init__.py:

_PROVIDERS = {
    "openai":    OpenAIProvider,
    "anthropic": AnthropicProvider,
    "gemini":    GeminiProvider,
    "ollama":    OllamaProvider,
}

LLM(provider="anthropic") calls _PROVIDERS["anthropic"](model=..., api_key=..., ...).

Lazy imports: Every provider SDK import (import openai, import anthropic, etc.) happens inside method bodies. Missing optional packages only raise errors when that specific provider is used, not when the library is imported. This means installing only anthropic still lets you import pywrapai without errors.

All four providers implement four abstract methods: chat(), stream(), achat(), and astream(). The async methods use each provider’s native async client (AsyncOpenAI, AsyncAnthropic, client.aio for Gemini, httpx.AsyncClient for Ollama) — not a thread pool wrapping the sync call.

Native JSON mode (Phase 1.3): Providers that support structured JSON output override a fifth optional method, chat_structured(messages). Returning an LLMResponse means native mode was used; returning None signals to LLM.chat_structured() that it should fall back to schema-injection + chat(). Raising an exception propagates immediately — it is not silently converted into a fallback.

Provider	chat_structured()	Mechanism
openai	overrides	response_format={"type": "json_object"}
gemini	overrides	response_mime_type="application/json"
ollama	overrides	"format": "json" in payload
anthropic	inherits base → None	prompt injection fallback

Streaming token capture: After a stream finishes, the provider stores the final token counts in self._last_stream_usage. LLM.stream() and LLM.astream() read and clear this value after the generator is exhausted, then record usage in TokenTracker.

---

Cache Key Architecture

The cache key is a SHA-256 hash of seven call parameters serialised as JSON:

SHA-256({
    "model":         "claude-haiku-4-5-20251001",
    "messages":      [...],           # full message list including role and content
    "temperature":   0.7,
    "max_tokens":    null,
    "system_prompt": "You are...",
    "tools":         [...] | null,    # neutral schema — added Phase 1.6
    "tool_choice":   "auto" | null,   # added Phase 1.6
})

tools and tool_choice were added in Phase 1.6 to ensure a tool-enabled call and a plain text call with the same prompt are never served each other’s cached response. All seven components must match exactly for a cache hit.

---

Token Tracking Architecture

LLM._do_chat(messages, tools, tool_choice)
  │
  ├── calls provider.chat(messages, tools, tool_choice)
  │   └── returns LLMResponse with usage.input_tokens, usage.output_tokens
  │       (token counts come from the API response — never estimated)
  │
  └── calls LLM._record_usage(response)
      └── TokenTracker.record(model, input_tokens, output_tokens, kind="chat")
          └── appends CallRecord(kind="chat") to internal list
              └── cost = get_cost(model, input_tokens, output_tokens)
                  └── looks up in PRICING table (USD per 1M tokens)

Embedder.embed(texts)                           ← Phase 2
  │
  ├── calls provider embed API
  │   └── token count from API response
  │       (Gemini/Ollama: 4-chars-per-token heuristic because embed APIs do not return counts)
  │
  └── TokenTracker.record(model, in_tokens, 0, kind="embedding")
      └── appends CallRecord(kind="embedding") to internal list
          └── cost looked up in embedding section of PRICING table

CallRecord.kind values:

• "chat" — default, used by every LLM.chat() call
• "embedding" — used by every Embedder.embed() call

TokenTracker.summary() breaks the totals down by kind:

Total calls   : 12  (10 chat + 2 embedding)
Total tokens  : 8,432  (input: 7,891  output: 512  embedding: 29)
Total cost    : $0.001247

A persistence layer built on top of GraphLLM can read the last entry from TokenTracker to get the cost of each call — get_cost(response.model, response.usage.input_tokens, response.usage.output_tokens). Embedding token records accumulate separately in the Embedder.tokens tracker and can be read by the application at any time.

---

File Structure

LLM-Library-Python/
│
├── pywrapai/                        ← Library 1
│   ├── __init__.py                  ← public exports
│   ├── llm.py                       ← LLM class (main entry point)
│   ├── embedder.py                  ← Embedder class (Phase 2; openai/gemini/ollama)
│   ├── core/
│   │   └── base.py                  ← Role, Message, TokenUsage, LLMResponse, ToolCall, ToolResult
│   ├── providers/
│   │   ├── __init__.py              ← get_provider() factory
│   │   ├── base.py                  ← BaseLLM abstract class
│   │   ├── openai.py
│   │   ├── anthropic.py
│   │   ├── gemini.py
│   │   └── ollama.py
│   ├── tokens/
│   │   ├── pricing.py               ← hardcoded prices + load_user_pricing + get_cost (chat + embedding models)
│   │   └── tracker.py               ← CallRecord (kind field), TokenTracker
│   ├── prompts/
│   │   └── template.py              ← PromptTemplate
│   ├── cache/
│   │   └── store.py                 ← BaseCache (abstract) + ResponseCache (in-memory LRU)
│   └── exceptions.py                ← PyWrapAIError hierarchy (AuthError, RateLimitError, …)
│
├── pywrapai_graph/                  ← Library 2
│   ├── __init__.py
│   ├── graph.py                     ← GraphLLM (with retriever=, rag_k=, rag= parameter)
│   ├── agent.py                     ← Agent, AgentConfig, AgentStep, AgentResult, AgentMaxTurnsError (Phase 3)
│   ├── prompts/
│   │   └── registry.py              ← PromptRegistry
│   ├── tools/
│   │   ├── __init__.py              ← re-exports Tool, ToolRegistry, tool
│   │   └── tool.py                  ← Tool dataclass, ToolRegistry, @tool decorator
│   └── rag/                         ← Phase 2 RAG pipeline
│       ├── __init__.py              ← re-exports all RAG public types + all store classes
│       ├── types.py                 ← Document, Chunk, ScoredChunk dataclasses
│       ├── chunker.py               ← Chunker (paragraph → sentence → hard-cut priority)
│       ├── store.py                 ← BaseVectorStore (ABC), InMemoryVectorStore (numpy L2-norm)
│       ├── retriever.py             ← BaseRetriever (ABC), VectorRetriever, retrieval_tool()
│       └── stores/                  ← Phase 2.1-2.5 persistent backends
│           ├── __init__.py          ← re-exports all five concrete backends
│           ├── sqlite.py            ← SQLiteVectorStore (stdlib sqlite3, no extra deps)
│           ├── chroma.py            ← ChromaVectorStore (pip install chromadb)
│           ├── faiss.py             ← FAISSVectorStore  (pip install faiss-cpu)
│           ├── pgvector.py          ← PGVectorStore     (pip install psycopg2-binary pgvector)
│           └── qdrant.py            ← QdrantVectorStore (pip install qdrant-client)
│
├── system_prompts_registry/         ← prompt files (user-managed)
│   ├── support-agent-v1.txt
│   └── support-agent-v2.txt
│
├── Gui-Application/                 ← reference implementation
│   ├── app.py                       ← PyQt5 chatbot (uses PyWrapAI + PyWrapAI-Graph)
│   └── app_db.py                    ← application-level user/session DB
│
├── test-agent-setup/                ← Phase 3 live test application (Flask + SQLite retail demo)
│   ├── seed_db.py                   ← creates store.db (33 products, 40 customers, 500 sales)
│   ├── db.py                        ← get_db() context manager; sqlite3.Row factory
│   ├── agent_setup.py               ← build_agent() + _log_step on_step callback
│   ├── app.py                       ← Flask server: POST /chat → agent.run()
│   ├── static/index.html            ← split GUI: chat (left) + tool execution log (right)
│   ├── .env                         ← local API keys (never committed)
│   ├── requirements.txt             ← flask, openai, python-dotenv
│   └── tools/
│       ├── search_products.py       ← @tool: name + supplier search, stock flags
│       ├── get_sales_stats.py       ← @tool: revenue, units, category breakdown by date range
│       ├── get_top_sellers.py       ← @tool: ranked by revenue or units sold
│       └── get_stock_status.py      ← @tool: inventory value, low-stock alerts (≤10 units)
│
├── Test-Example/                    ← test suite (397 PASS + 9 SKIP)
│   ├── test_tools.py                ← Phase 1.6 + closing tasks: 22 unit tests, no live API calls
│   ├── test_tools_live.py           ← Phase 1.6: 5 live tests, requires ANTHROPIC_API_KEY
│   ├── test_graph.py                ← GraphLLM prompt versioning: 6 tests
│   ├── test_graph_integration.py    ← GraphLLM end-to-end: 6 tests
│   ├── test_history.py              ← Conversation history injection: 6 tests
│   ├── test_anthropic.py            ← Anthropic provider live: 7 tests
│   ├── test_rag.py                  ← Phase 2 RAG pipeline: 11 unit tests + 2 SKIP (live embed)
│   ├── test_vector_stores.py        ← Phase 2.1-2.5: 37 PASS + 0 FAIL + 7 SKIP (PG live tests)
│   └── test_agent.py                ← Phase 3 Agent ReAct loop: 229 unit tests, no API calls
│
└── docs/                            ← documentation
    ├── tutorial.md                  ← step-by-step guide through PyWrapAI + PyWrapAI-Graph + RAG + stores
    ├── pywrapai.md                  ← Library 1 reference
    ├── pywrapai-graph.md            ← Library 2 reference (includes RAG + all vector store backends)
    ├── architecture.md              ← this file — system overview, message flow, file structure
    ├── vector_db_logic.md           ← vector store backend deep-dive (all 6 backends, test results)
    ├── rag-chatbot-app.md           ← RAG-Test-GUI-App reference (setup, architecture, debugging)
    ├── lib-finetuning.md            ← remaining tasks and planned improvements
    ├── phase0-completion.md         ← Phase 0 completion notes
    └── phase1-progress.md           ← Phase 1 + Phase 2 (including 2.1-2.5) progress report

---

Provider Tool Wire Format (Phase 1.6)

All four providers implement the same BaseLLM.chat(tools=, tool_choice=) abstract signature. Each converts the neutral format to its own wire format internally. Application code always passes registry.to_schema() and never needs to know which format each provider uses.

Provider	Tool format	Tool choice	Response tool calls	Arguments format
OpenAI	[{"type":"function","function":{...}}]	{"type":"function","function":{"name":"..."}} for named	choices[0].message.tool_calls	JSON string in request; parsed to dict on response
Anthropic	[{"name":...,"input_schema":{...}}]	{"type":"any"} for "required"; {"type":"tool","name":...} for named	content blocks of type "tool_use"	dict directly; NO "tool" role on wire — Role.TOOL becomes user-role with tool_result block
Gemini	FunctionDeclaration objects via SDK	tool_config parameter	candidates[0].content.parts with function_call	dict directly; response.text is None on tool calls — must walk parts; no call IDs, uuid4() fills the gap; additionalProperties/default/title/$schema/$defs keys stripped from schema
Ollama	Same as OpenAI (OpenAI-compatible wire)	Not supported — emits warning and ignores	message.tool_calls	May be JSON string or dict — both handled

The key insight: only the provider files (openai.py, anthropic.py, gemini.py, ollama.py) know about these differences. LLM, GraphLLM, and all test code use the neutral format exclusively.

---

Design Philosophy

One library, one job. This is not a single engine that does everything (like LangChain). Each library has a single, clearly stated responsibility. You can use PyWrapAI alone in a data pipeline. You add PyWrapAI-Graph only when you need message assembly, RAG, or tool calling.

No magic. Every behaviour is explicit. Tokens are counted by the API, not estimated. Cost calculation uses a visible, overridable price table. History is a plain list your application controls — nothing is loaded or saved on your behalf without you asking for it.

Production first. The design targets real deployed applications serving thousands of users simultaneously across time zones — not local scripts. Caching, retry, and fallback are built in from the start rather than bolted on later.