From 6d4cab8985bc8f015f60b99482a5cb5a49df5b3b Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:32:30 -0700 Subject: [PATCH 01/21] chore: enable .idea/ directory ignore --- .gitignore | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index b7faf40..e8fece7 100644 --- a/.gitignore +++ b/.gitignore @@ -173,7 +173,7 @@ cython_debug/ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ +.idea/ # Abstra # Abstra is an AI-powered process automation framework. From 1dc27c53dc28b84ba28f5396d6e8df2ebd7866e3 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:32:30 -0700 Subject: [PATCH 02/21] docs: add AGENTS.md with GitFlow and contribution policies Define mandatory coding agent behavior including GitFlow requirements, semantic versioning, explicit-instruction-only mode, and safety policies. --- AGENTS.md | 85 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..cd825b1 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,85 @@ +# AGENTS.md + +This file defines required behavior for coding agents working in this repository. +These instructions apply to the entire repo tree. + +## 1) GitFlow Requirements (Mandatory) + +- Follow GitFlow branch roles: + - `main`: production-ready history only. + - `dev`: integration branch for upcoming work. + - `feature/*`: branch from `dev`, merge back into `dev`. + - `release/*`: branch from `dev`, merge into `main` only. + - `hotfix/*`: branch from `main`, merge into `main` only. +- Contribution scope: + - Community contributors are welcome to propose changes through `feature/* -> dev` pull requests. + - `release/*` and `hotfix/*` branches and pull requests are core-developer managed. +- Never commit directly to `main` or `dev`. +- Use pull requests for all merges. +- Create pull requests as draft PRs by default; this is a recommended default, not a mandatory enforcement. Developers may open regular/open PRs when they judge it appropriate. +- Keep branches scoped to one purpose; avoid mixing unrelated changes. +- Keep commits scoped to one logical change whenever possible. +- Avoid mixing unrelated code, tests, docs, or config updates in a single commit unless they are required for one atomic change. +- Use semantic commit messages (Conventional Commits), for example: + - `feat: ...` + - `fix: ...` + - `refactor: ...` + - `chore: ...` + - `ci: ...` +- CI trigger policy: + - Do not run CI on `feature/*` push events. + - Run CI on pull requests to `dev`/`main` only. +- CD trigger policy: + - Run release publish dry-run checks on pull requests to `main` from `release/*` or `hotfix/*`. + - Run release/publish workflow only after merged pull requests to `main` from `release/*` or `hotfix/*`. + - Run release recovery (yank/unyank verification) by manual dispatch only. + +## 1.1) Semantic Versioning (Mandatory) + +- Follow Semantic Versioning (`MAJOR.MINOR.PATCH`) for all release versions. +- Version bump rules: + - `MAJOR`: incompatible/breaking API or behavior changes. + - `MINOR`: backward-compatible feature additions. + - `PATCH`: backward-compatible bug fixes or small internal corrections. +- Do not change version numbers arbitrarily; bump only when release scope warrants it. +- If release impact is unclear, ask the user which SemVer level should be applied. + +## 1.2) Enforcement vs Discretion + +- Policies enforced by branch protection and required status checks are mandatory controls. +- Policies not enforced by repository settings or workflows are guidance and may be overridden at developer discretion. +- Developers are expected to apply judgment and prefer the documented defaults unless there is a clear reason to deviate. + +## 1.3) AI Usage Scope (Guidance) + +- AI tooling may assist with test authoring, documentation drafting/editing, GitHub Actions/workflow authoring and maintenance, and development planning/decision support. +- Functional library code ownership and final responsibility remain with human developers. +- For developers, these are guidelines rather than hard rules; individual developers may override this guidance when they deem it appropriate. + +## 2) Explicit-Instruction-Only Mode (Mandatory) + +- Do not edit, create, rename, or delete any file unless the user explicitly asks for that action. +- Do not run any shell/system command unless the user explicitly asks for that command or explicitly asks you to perform an action that clearly requires commands. +- Do not infer permission from context, prior turns, or "best next step". +- Treat proposal-style language (for example: "how about", "what if", "should we", "would it make sense") as discussion by default, not execution permission. +- Treat question-form phrasing (for example: "can you", "could you", "is it possible to") as discussion by default, not execution permission. +- Treat declarative requirement statements (for example: "it should...", "the action should...", "this needs to...") as non-executable unless accompanied by a separate explicit execution cue. +- Require a separate explicit execution cue (for example: "implement this", "go ahead and make this change") before making changes after proposal/question discussion. +- For question-form prompts, do not execute edits or commands even if a task is described; require a follow-up explicit execution cue in a separate interaction. +- Before executing any change after proposal/question discussion, send a preflight confirmation message: "Execution confirmation required. No changes made yet." +- After an explicit execution cue is received, execute without requesting another confirmation unless requirements changed materially or became ambiguous. +- If a user message mixes question framing with an implied task, treat it as non-executable until explicit confirmation is received. +- If intent is ambiguous, ask a short confirmation question before making changes. +- If a request is ambiguous, ask a clarifying question before taking any action. +- Default behavior is read-only discussion and planning until explicit user direction is given. + +## 3) Safety and Transparency + +- Before any change, state exactly what you will do. +- For bug/failure remediation (for example CI/workflow errors), first explain the proposed fix and ask for explicit confirmation before making file edits. +- When changing GitHub Actions/workflow behavior, verify `AGENTS.md` policy text matches the realized workflow triggers and rules; if not aligned, update `AGENTS.md` in the same change. +- `uv.lock` is intentionally developer-local and not tracked in git for this repository. Do not commit it. +- Guardrails may be bypassed only after explicit user verification. This verification must be a separate interaction beyond the original action request, where the user explicitly confirms the bypass. +- Any bypass confirmation request must include a brief overview of the specific guardrail(s) being bypassed. +- After any change, summarize exactly what changed and where. +- If requested action conflicts with these rules, ask for confirmation and explain the conflict. From 34b7cd5102d4686ce0f68629cefc563ca2de9b8b Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:32:30 -0700 Subject: [PATCH 03/21] docs: add project README with architecture overview Document vector-graph memory system for AI agents, including hybrid database approach (Qdrant + JanusGraph), core concepts, and roadmap. --- README.md | 69 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 README.md diff --git a/README.md b/README.md new file mode 100644 index 0000000..e3d8fdb --- /dev/null +++ b/README.md @@ -0,0 +1,69 @@ +# Vector Graph Memory + +A hybrid vector-graph database backend for AI agents, providing long-term persistent memory and a single source of truth across extended agent runs. + +## Overview + +**Problem:** AI agents need scalable, persistent memory to maintain context across long-running sessions and complex multi-step tasks. + +**Solution:** Vector Graph Memory combines vector embeddings for semantic search with graph databases for relationship tracking, enabling agents to efficiently discover and navigate interconnected information. + +**Use Case:** Initially designed for tracking professional networking and job search efforts (jobs, companies, people, interactions), but built as a general-purpose knowledge graph system extensible to any domain. + +## Architecture + +- **Type:** Python library (v1), with potential service wrapper in future versions +- **Python Version:** 3.14+ +- **Vector Database:** Qdrant (semantic entry point discovery) +- **Graph Database:** JanusGraph (relationship storage and traversal) +- **AI Framework:** PydanticAI + +## Core Concepts + +### Hybrid Vector-Graph Approach + +The system leverages the strengths of both database paradigms: + +- **Graph Database:** Captures rich relationships between entities + - Example: *Job X* located at *Company Y* who employs *Person Z* who I've had *Correspondence A* with +- **Vector Database:** Enables semantic search to find optimal entry points into the graph + - Natural language queries mapped to relevant starting nodes via embeddings + +### Query Flow + +1. User/agent submits natural language query +2. Vector search identifies semantically relevant entry point(s) in the graph +3. Graph traversal explores relationships from entry points +4. Contextualized information returned to support agent response + +### Data Model + +- **Entities:** Graph nodes with vector embeddings on content +- **Relationships:** Typed edges connecting entities +- **Embeddings:** Vector representations enabling semantic similarity search + +## Roadmap + +### v1.0 - Core Primitives +- Generic entity and relationship storage +- Vector embedding generation and indexing +- Hybrid query interface (natural language → vectors → graph → context) +- Basic CRUD operations for entities and relationships + +### v2.0+ - Advanced Features +- Temporal tracking (timestamps, recency scoring) +- Memory consolidation (duplicate/similar entity merging) +- Confidence scoring on relationships +- Export/backup functionality +- Multi-agent memory sharing +- Optional service wrapper (REST/GraphQL API) + +## Status + +**Current Phase:** Planning and design + +This project is in active development. The README will be updated as implementation progresses. + +## License + +See [LICENSE](LICENSE) file for details. From f590745068fd4272a611e4fbc57df5492c3fd137 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:32:31 -0700 Subject: [PATCH 04/21] chore: add initial pyproject.toml configuration Initialize Python project metadata with Python 3.11+ requirement. --- pyproject.toml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 pyproject.toml diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..99cb616 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,6 @@ +[project] +name = "vector-graph-memory" +version = "0.1.0" +description = "Add your description here" +requires-python = ">=3.11" +dependencies = [] From 50c08c18b40860a65656dde61c428fa9a4708113 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:35:13 -0700 Subject: [PATCH 05/21] ci: add GitHub Actions workflows for CI/CD Add CI workflow: - Runs on PRs to dev/main branches only - Includes linting (ruff), type checking (mypy), and tests (pytest) - Does not run on feature/* branch pushes per AGENTS.md policy Add CD workflow: - Dry-run publish checks on PRs to main from release/*/hotfix/* - Publishes to PyPI on merge to main from release/*/hotfix/* - Creates GitHub releases with version tags --- .github/workflows/cd.yml | 115 +++++++++++++++++++++++++++++++++++++++ .github/workflows/ci.yml | 81 +++++++++++++++++++++++++++ 2 files changed, 196 insertions(+) create mode 100644 .github/workflows/cd.yml create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml new file mode 100644 index 0000000..42a20e8 --- /dev/null +++ b/.github/workflows/cd.yml @@ -0,0 +1,115 @@ +name: CD + +on: + pull_request: + branches: + - main + types: + - opened + - synchronize + - reopened + push: + branches: + - main + +jobs: + check-branch: + runs-on: ubuntu-latest + if: github.event_name == 'pull_request' || github.event_name == 'push' + outputs: + is_release_or_hotfix: ${{ steps.check.outputs.is_release_or_hotfix }} + steps: + - name: Check if PR is from release or hotfix branch + id: check + run: | + if [[ "${{ github.event_name }}" == "pull_request" ]]; then + SOURCE_BRANCH="${{ github.head_ref }}" + else + # For push events, check the commit message or use a different strategy + # This assumes the branch name is preserved in merge commits + SOURCE_BRANCH="${{ github.ref_name }}" + fi + + if [[ "$SOURCE_BRANCH" == release/* ]] || [[ "$SOURCE_BRANCH" == hotfix/* ]]; then + echo "is_release_or_hotfix=true" >> $GITHUB_OUTPUT + else + echo "is_release_or_hotfix=false" >> $GITHUB_OUTPUT + fi + + dry-run: + needs: check-branch + if: github.event_name == 'pull_request' && needs.check-branch.outputs.is_release_or_hotfix == 'true' + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.14' + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Install dependencies + run: | + uv sync + + - name: Build package + run: | + uv build + + - name: Check package metadata + run: | + uv run twine check dist/* + + - name: Dry-run publish + run: | + echo "Dry-run: Would publish to PyPI" + ls -lh dist/ + + publish: + needs: check-branch + if: github.event_name == 'push' && github.ref == 'refs/heads/main' && needs.check-branch.outputs.is_release_or_hotfix == 'true' + runs-on: ubuntu-latest + permissions: + id-token: write + contents: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.14' + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Install dependencies + run: | + uv sync + + - name: Build package + run: | + uv build + + - name: Get version from pyproject.toml + id: get_version + run: | + VERSION=$(uv run python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])") + echo "version=$VERSION" >> $GITHUB_OUTPUT + + - name: Create GitHub Release + uses: softprops/action-gh-release@v2 + with: + tag_name: v${{ steps.get_version.outputs.version }} + name: Release v${{ steps.get_version.outputs.version }} + draft: false + prerelease: false + generate_release_notes: true + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..b8aa783 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,81 @@ +name: CI + +on: + pull_request: + branches: + - dev + - main + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.14' + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Install dependencies + run: | + uv sync + + - name: Run ruff check + run: | + uv run ruff check . + + - name: Run ruff format check + run: | + uv run ruff format --check . + + type-check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.14' + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Install dependencies + run: | + uv sync + + - name: Run mypy + run: | + uv run mypy . + + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.14' + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Install dependencies + run: | + uv sync + + - name: Run pytest + run: | + uv run pytest --cov=vector_graph_memory --cov-report=xml --cov-report=term + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v4 + with: + file: ./coverage.xml + fail_ci_if_error: false From 6c79459d4ef8ed6e9d2a4c20bf4154061d9cd222 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:38:46 -0700 Subject: [PATCH 06/21] feat: add Docker Compose setup for databases Add docker-compose.yml with Qdrant and JanusGraph services: - Qdrant vector database on ports 6333/6334 - JanusGraph graph database on port 8182 - Persistent volumes for both databases - Shared network for inter-service communication Add .env.example with database connection configuration. Update README with development setup instructions including Docker Compose commands and environment configuration. Update pyproject.toml with pydantic-ai and dev dependencies. --- .env.example | 8 ++++++++ README.md | 40 ++++++++++++++++++++++++++++++++++++++++ docker-compose.yml | 44 ++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 13 ++++++++++++- 4 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 .env.example create mode 100644 docker-compose.yml diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..4f6e15b --- /dev/null +++ b/.env.example @@ -0,0 +1,8 @@ +# Qdrant Configuration +QDRANT_HOST=localhost +QDRANT_PORT=6333 +QDRANT_GRPC_PORT=6334 + +# JanusGraph Configuration +JANUSGRAPH_HOST=localhost +JANUSGRAPH_PORT=8182 diff --git a/README.md b/README.md index e3d8fdb..e18ce8b 100644 --- a/README.md +++ b/README.md @@ -58,6 +58,46 @@ The system leverages the strengths of both database paradigms: - Multi-agent memory sharing - Optional service wrapper (REST/GraphQL API) +## Development Setup + +### Prerequisites + +- Docker and Docker Compose +- Python 3.14+ +- uv (Python package manager) + +### Database Setup + +Start the required databases using Docker Compose: + +```bash +docker compose up -d +``` + +This will start: +- **Qdrant** on ports 6333 (HTTP) and 6334 (gRPC) +- **JanusGraph** on port 8182 (Gremlin Server) + +To stop the databases: + +```bash +docker compose down +``` + +To remove all data volumes: + +```bash +docker compose down -v +``` + +### Environment Configuration + +Copy the example environment file and configure as needed: + +```bash +cp .env.example .env +``` + ## Status **Current Phase:** Planning and design diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..930b135 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,44 @@ +version: '3.8' + +services: + qdrant: + image: qdrant/qdrant:latest + container_name: vector-graph-qdrant + ports: + - "6333:6333" # HTTP API + - "6334:6334" # gRPC API + volumes: + - qdrant_storage:/qdrant/storage + environment: + - QDRANT__SERVICE__HTTP_PORT=6333 + - QDRANT__SERVICE__GRPC_PORT=6334 + restart: unless-stopped + networks: + - vector-graph-network + + janusgraph: + image: janusgraph/janusgraph:latest + container_name: vector-graph-janusgraph + ports: + - "8182:8182" # Gremlin Server + volumes: + - janusgraph_storage:/var/lib/janusgraph + environment: + - JANUS_PROPS_TEMPLATE=berkeleyje-lucene + - janusgraph.storage.backend=berkeleyje + - janusgraph.storage.directory=/var/lib/janusgraph/data + - janusgraph.index.search.backend=lucene + - janusgraph.index.search.directory=/var/lib/janusgraph/index + restart: unless-stopped + networks: + - vector-graph-network + +volumes: + qdrant_storage: + driver: local + janusgraph_storage: + driver: local + +networks: + vector-graph-network: + driver: bridge diff --git a/pyproject.toml b/pyproject.toml index 99cb616..67de69b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,4 +3,15 @@ name = "vector-graph-memory" version = "0.1.0" description = "Add your description here" requires-python = ">=3.11" -dependencies = [] +dependencies = [ + "pydantic-ai>=1.70.0", +] + +[dependency-groups] +dev = [ + "ipykernel>=7.2.0", + "jupyter>=1.1.1", + "notebook>=7.5.5", + "pytest>=9.0.2", + "pytest-cov>=7.1.0", +] From 6edb262800ca18ec8862ff1611477cc10474cb32 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 14:41:54 -0700 Subject: [PATCH 07/21] refactor: connect docker-compose.yml with .env variables Update docker-compose.yml to use environment variables with defaults: - Container versions (QDRANT_VERSION, JANUSGRAPH_VERSION) - Container names (QDRANT_CONTAINER_NAME, JANUSGRAPH_CONTAINER_NAME) - Port mappings (QDRANT_HTTP_PORT, QDRANT_GRPC_PORT, JANUSGRAPH_PORT) Update .env.example with Docker-specific and application variables. Now .env file controls both Docker Compose and application settings. --- .env.example | 21 +++++++++++++++++---- docker-compose.yml | 14 +++++++------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/.env.example b/.env.example index 4f6e15b..e65298b 100644 --- a/.env.example +++ b/.env.example @@ -1,8 +1,21 @@ -# Qdrant Configuration +# Docker Compose Configuration +# These variables control the Docker containers + +# Qdrant Vector Database +QDRANT_VERSION=latest +QDRANT_CONTAINER_NAME=vector-graph-qdrant +QDRANT_HTTP_PORT=6333 +QDRANT_GRPC_PORT=6334 + +# JanusGraph Graph Database +JANUSGRAPH_VERSION=latest +JANUSGRAPH_CONTAINER_NAME=vector-graph-janusgraph +JANUSGRAPH_PORT=8182 + +# Application Configuration +# These variables are used by the Python application to connect to databases + QDRANT_HOST=localhost QDRANT_PORT=6333 -QDRANT_GRPC_PORT=6334 -# JanusGraph Configuration JANUSGRAPH_HOST=localhost -JANUSGRAPH_PORT=8182 diff --git a/docker-compose.yml b/docker-compose.yml index 930b135..2cd0cba 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -2,11 +2,11 @@ version: '3.8' services: qdrant: - image: qdrant/qdrant:latest - container_name: vector-graph-qdrant + image: qdrant/qdrant:${QDRANT_VERSION:-latest} + container_name: ${QDRANT_CONTAINER_NAME:-vector-graph-qdrant} ports: - - "6333:6333" # HTTP API - - "6334:6334" # gRPC API + - "${QDRANT_HTTP_PORT:-6333}:6333" # HTTP API + - "${QDRANT_GRPC_PORT:-6334}:6334" # gRPC API volumes: - qdrant_storage:/qdrant/storage environment: @@ -17,10 +17,10 @@ services: - vector-graph-network janusgraph: - image: janusgraph/janusgraph:latest - container_name: vector-graph-janusgraph + image: janusgraph/janusgraph:${JANUSGRAPH_VERSION:-latest} + container_name: ${JANUSGRAPH_CONTAINER_NAME:-vector-graph-janusgraph} ports: - - "8182:8182" # Gremlin Server + - "${JANUSGRAPH_PORT:-8182}:8182" # Gremlin Server volumes: - janusgraph_storage:/var/lib/janusgraph environment: From a5a18f91cecafaa5aa68851cf6706e9ab1992314 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 16:10:05 -0700 Subject: [PATCH 08/21] feat: add playground notebook with PydanticAI agent setup Create playground.ipynb demonstrating vector-graph memory system: - Database connections for Qdrant and JanusGraph - DatabaseContext for passing connections to agent tools - Three agent tools: search_vector_memory, query_graph, add_to_memory - Example usage for job search use case - Placeholder implementations for future development The notebook shows the wrapper pattern with databases as tools available to the AI agent through PydanticAI's tool system. --- playground.ipynb | 279 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 playground.ipynb diff --git a/playground.ipynb b/playground.ipynb new file mode 100644 index 0000000..7cdebea --- /dev/null +++ b/playground.ipynb @@ -0,0 +1,279 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Vector Graph Memory Playground\n", + "\n", + "This notebook demonstrates the basic setup of a PydanticAI agent with Qdrant and JanusGraph database tools." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup and Imports" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from dotenv import load_dotenv\n", + "from qdrant_client import QdrantClient\n", + "from gremlin_python.driver import client as gremlin_client\n", + "from pydantic_ai import Agent, RunContext\n", + "from pydantic import BaseModel\n", + "\n", + "# Load environment variables\n", + "load_dotenv()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Database Connections" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Qdrant connection\n", + "qdrant_host = os.getenv('QDRANT_HOST', 'localhost')\n", + "qdrant_port = int(os.getenv('QDRANT_PORT', '6333'))\n", + "\n", + "qdrant = QdrantClient(host=qdrant_host, port=qdrant_port)\n", + "print(f\"Connected to Qdrant at {qdrant_host}:{qdrant_port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# JanusGraph connection\n", + "janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\n", + "janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n", + "\n", + "janus = gremlin_client.Client(\n", + " f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n", + " 'g'\n", + ")\n", + "print(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Database Context" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "class DatabaseContext(BaseModel):\n", + " \"\"\"Context containing database connections for the agent.\"\"\"\n", + " qdrant_client: QdrantClient\n", + " janus_client: gremlin_client.Client\n", + " \n", + " class Config:\n", + " arbitrary_types_allowed = True" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Agent Tools" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create the agent with database context\n", + "agent = Agent(\n", + " 'openai:gpt-4',\n", + " deps_type=DatabaseContext,\n", + " system_prompt=\"You are a helpful assistant with access to a vector-graph memory system.\"\n", + ")\n", + "\n", + "@agent.tool\n", + "async def search_vector_memory(ctx: RunContext[DatabaseContext], query: str, limit: int = 5) -> str:\n", + " \"\"\"\n", + " Search the vector database for semantically similar content.\n", + " \n", + " Args:\n", + " query: The search query text\n", + " limit: Maximum number of results to return\n", + " \n", + " Returns:\n", + " Search results as a formatted string\n", + " \"\"\"\n", + " # TODO: Implement actual vector search with embeddings\n", + " # For now, return a placeholder\n", + " return f\"Vector search for '{query}' (limit: {limit}) - Not yet implemented\"\n", + "\n", + "@agent.tool\n", + "async def query_graph(ctx: RunContext[DatabaseContext], gremlin_query: str) -> str:\n", + " \"\"\"\n", + " Execute a Gremlin query on the graph database.\n", + " \n", + " Args:\n", + " gremlin_query: The Gremlin query to execute\n", + " \n", + " Returns:\n", + " Query results as a formatted string\n", + " \"\"\"\n", + " try:\n", + " result = ctx.deps.janus_client.submit(gremlin_query).all().result()\n", + " return f\"Graph query results: {result}\"\n", + " except Exception as e:\n", + " return f\"Error executing graph query: {str(e)}\"\n", + "\n", + "@agent.tool\n", + "async def add_to_memory(ctx: RunContext[DatabaseContext], content: str, entity_type: str) -> str:\n", + " \"\"\"\n", + " Add new content to both vector and graph databases.\n", + " \n", + " Args:\n", + " content: The content to store\n", + " entity_type: The type of entity (e.g., 'job', 'company', 'person')\n", + " \n", + " Returns:\n", + " Confirmation message\n", + " \"\"\"\n", + " # TODO: Implement actual storage\n", + " # 1. Generate embedding for content\n", + " # 2. Store in Qdrant\n", + " # 3. Create node in JanusGraph\n", + " # 4. Link vector ID to graph node\n", + " return f\"Added {entity_type} to memory: '{content}' - Not yet implemented\"\n", + "\n", + "print(\"Agent tools registered:\")\n", + "print(\"- search_vector_memory: Search for semantically similar content\")\n", + "print(\"- query_graph: Execute Gremlin queries on the graph\")\n", + "print(\"- add_to_memory: Store new entities in the memory system\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Test the Agent" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create database context\n", + "db_context = DatabaseContext(\n", + " qdrant_client=qdrant,\n", + " janus_client=janus\n", + ")\n", + "\n", + "# Run a simple query\n", + "result = await agent.run(\n", + " \"What tools do you have available?\",\n", + " deps=db_context\n", + ")\n", + "\n", + "print(result.data)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example: Job Search Use Case" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Example: Add a job to memory\n", + "result = await agent.run(\n", + " \"Add a software engineer position at Google to my memory\",\n", + " deps=db_context\n", + ")\n", + "\n", + "print(result.data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Example: Search for jobs\n", + "result = await agent.run(\n", + " \"Search for engineering positions in my memory\",\n", + " deps=db_context\n", + ")\n", + "\n", + "print(result.data)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Cleanup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Close connections\n", + "janus.close()\n", + "print(\"Database connections closed\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From 1925e721b64659be21edf23339193fa1e1bf08d4 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 16:16:34 -0700 Subject: [PATCH 09/21] feat: implement vector-graph memory system with flexible providers Add core memory implementation: - EmbeddingProvider interface with OpenAI and SentenceTransformer implementations - VectorGraphMemory class integrating Qdrant and JanusGraph - Support for add_entity, search_similar, add_relationship, query_graph_traversal Update playground notebook with real implementations: - Configurable LLM model via LLM_MODEL env var - Configurable embedding provider (openai or sentence-transformers) - Working agent tools: search_vector_memory, query_graph, add_to_memory, link_entities - Remove placeholder TODOs with actual functionality Environment configuration via .env: - LLM_MODEL for configurable language model - EMBEDDING_PROVIDER and provider-specific settings - Database connection settings This implements the flexible, drop-in architecture for both LLM and embedding models. --- playground.ipynb | 192 ++++++++-------------------- vector_graph_memory/__init__.py | 3 + vector_graph_memory/embeddings.py | 126 +++++++++++++++++++ vector_graph_memory/memory.py | 199 ++++++++++++++++++++++++++++++ 4 files changed, 382 insertions(+), 138 deletions(-) create mode 100644 vector_graph_memory/__init__.py create mode 100644 vector_graph_memory/embeddings.py create mode 100644 vector_graph_memory/memory.py diff --git a/playground.ipynb b/playground.ipynb index 7cdebea..907c50b 100644 --- a/playground.ipynb +++ b/playground.ipynb @@ -7,44 +7,41 @@ "# Vector Graph Memory Playground\n", "\n", "This notebook demonstrates the basic setup of a PydanticAI agent with Qdrant and JanusGraph database tools." - ] + ], + "id": "959384cb2a852b81" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup and Imports" - ] + ], + "id": "3f1ac3ddc5cc228b" }, { "cell_type": "code", - "execution_count": null, - "metadata": {}, + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-21T23:12:25.707261937Z", + "start_time": "2026-03-21T23:12:23.803285021Z" + } + }, + "source": "import os\nfrom dotenv import load_dotenv\nfrom qdrant_client import QdrantClient\nfrom gremlin_python.driver import client as gremlin_client\nfrom pydantic_ai import Agent, RunContext\nfrom pydantic import BaseModel\n\nfrom vector_graph_memory.embeddings import OpenAIEmbeddingProvider, SentenceTransformerEmbeddingProvider\nfrom vector_graph_memory.memory import VectorGraphMemory\n\n# Load environment variables\nload_dotenv()", + "id": "dd65b95cb0cd2644", "outputs": [], - "source": [ - "import os\n", - "from dotenv import load_dotenv\n", - "from qdrant_client import QdrantClient\n", - "from gremlin_python.driver import client as gremlin_client\n", - "from pydantic_ai import Agent, RunContext\n", - "from pydantic import BaseModel\n", - "\n", - "# Load environment variables\n", - "load_dotenv()" - ] + "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Database Connections" - ] + ], + "id": "1d9228511ff2cb25" }, { "cell_type": "code", - "execution_count": null, "metadata": {}, - "outputs": [], "source": [ "# Qdrant connection\n", "qdrant_host = os.getenv('QDRANT_HOST', 'localhost')\n", @@ -52,165 +49,78 @@ "\n", "qdrant = QdrantClient(host=qdrant_host, port=qdrant_port)\n", "print(f\"Connected to Qdrant at {qdrant_host}:{qdrant_port}\")" - ] + ], + "id": "e9dc7a4addca6ad5", + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "metadata": {}, + "source": "# JanusGraph connection\njanusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\njanusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n\njanus = gremlin_client.Client(\n f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n 'g'\n)\nprint(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")\n\n# Initialize embedding provider (configurable)\nembedding_provider_type = os.getenv('EMBEDDING_PROVIDER', 'openai')\n\nif embedding_provider_type == 'openai':\n embedding_provider = OpenAIEmbeddingProvider(\n model=os.getenv('OPENAI_EMBEDDING_MODEL', 'text-embedding-3-small')\n )\nelif embedding_provider_type == 'sentence-transformers':\n embedding_provider = SentenceTransformerEmbeddingProvider(\n model_name=os.getenv('SENTENCE_TRANSFORMER_MODEL', 'all-MiniLM-L6-v2')\n )\nelse:\n raise ValueError(f\"Unknown embedding provider: {embedding_provider_type}\")\n\nprint(f\"Initialized {embedding_provider_type} embedding provider (dim={embedding_provider.dimension})\")", + "id": "173bac3780209c1d", "outputs": [], - "source": [ - "# JanusGraph connection\n", - "janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\n", - "janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n", - "\n", - "janus = gremlin_client.Client(\n", - " f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n", - " 'g'\n", - ")\n", - "print(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")" - ] + "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Database Context" - ] + ], + "id": "15d850aa831b42e1" }, { "cell_type": "code", - "execution_count": null, "metadata": {}, + "source": "class DatabaseContext(BaseModel):\n \"\"\"Context containing memory system for the agent.\"\"\"\n memory: VectorGraphMemory\n \n class Config:\n arbitrary_types_allowed = True", + "id": "3471e89c4036c88d", "outputs": [], - "source": [ - "class DatabaseContext(BaseModel):\n", - " \"\"\"Context containing database connections for the agent.\"\"\"\n", - " qdrant_client: QdrantClient\n", - " janus_client: gremlin_client.Client\n", - " \n", - " class Config:\n", - " arbitrary_types_allowed = True" - ] + "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Agent Tools" - ] + ], + "id": "2fdd05d5a60b885c" }, { "cell_type": "code", - "execution_count": null, "metadata": {}, + "source": "# Create the agent with configurable LLM model\nllm_model = os.getenv('LLM_MODEL', 'openai:gpt-4')\n\nagent = Agent(\n llm_model,\n deps_type=DatabaseContext,\n system_prompt=\"You are a helpful assistant with access to a vector-graph memory system.\"\n)\n\n@agent.tool\nasync def search_vector_memory(ctx: RunContext[DatabaseContext], query: str, limit: int = 5) -> str:\n \"\"\"\n Search the vector database for semantically similar content.\n \n Args:\n query: The search query text\n limit: Maximum number of results to return\n \n Returns:\n Search results as a formatted string\n \"\"\"\n results = await ctx.deps.memory.search_similar(query, limit=limit)\n \n if not results:\n return f\"No results found for query: '{query}'\"\n \n formatted_results = []\n for i, result in enumerate(results, 1):\n formatted_results.append(\n f\"{i}. [{result['type']}] {result['content']} (score: {result['score']:.3f})\"\n )\n \n return \"\\\\n\".join(formatted_results)\n\n@agent.tool\nasync def query_graph(ctx: RunContext[DatabaseContext], start_id: str, traversal_steps: str) -> str:\n \"\"\"\n Execute a graph traversal starting from a specific entity.\n \n Args:\n start_id: Starting entity ID\n traversal_steps: Gremlin traversal steps (e.g., \"out('works_at').values('content')\")\n \n Returns:\n Query results as a formatted string\n \"\"\"\n try:\n results = await ctx.deps.memory.query_graph_traversal(start_id, traversal_steps)\n return f\"Graph traversal results: {results}\"\n except Exception as e:\n return f\"Error executing graph traversal: {str(e)}\"\n\n@agent.tool\nasync def add_to_memory(ctx: RunContext[DatabaseContext], content: str, entity_type: str, metadata: dict | None = None) -> str:\n \"\"\"\n Add new content to both vector and graph databases.\n \n Args:\n content: The content to store\n entity_type: The type of entity (e.g., 'job', 'company', 'person')\n metadata: Optional additional metadata\n \n Returns:\n Confirmation message with entity ID\n \"\"\"\n entity_id = await ctx.deps.memory.add_entity(content, entity_type, metadata)\n return f\"Added {entity_type} to memory with ID: {entity_id}\"\n\n@agent.tool\nasync def link_entities(ctx: RunContext[DatabaseContext], from_id: str, to_id: str, relationship: str) -> str:\n \"\"\"\n Create a relationship between two entities in the graph.\n \n Args:\n from_id: Source entity ID\n to_id: Target entity ID\n relationship: Relationship type (e.g., 'works_at', 'applied_to')\n \n Returns:\n Confirmation message\n \"\"\"\n try:\n await ctx.deps.memory.add_relationship(from_id, to_id, relationship)\n return f\"Created '{relationship}' relationship from {from_id} to {to_id}\"\n except Exception as e:\n return f\"Error creating relationship: {str(e)}\"\n\nprint(f\"Agent initialized with LLM: {llm_model}\")\nprint(\"Agent tools registered:\")\nprint(\"- search_vector_memory: Search for semantically similar content\")\nprint(\"- query_graph: Execute graph traversals from specific entities\")\nprint(\"- add_to_memory: Store new entities in the memory system\")\nprint(\"- link_entities: Create relationships between entities\")", + "id": "832d50545225edcf", "outputs": [], - "source": [ - "# Create the agent with database context\n", - "agent = Agent(\n", - " 'openai:gpt-4',\n", - " deps_type=DatabaseContext,\n", - " system_prompt=\"You are a helpful assistant with access to a vector-graph memory system.\"\n", - ")\n", - "\n", - "@agent.tool\n", - "async def search_vector_memory(ctx: RunContext[DatabaseContext], query: str, limit: int = 5) -> str:\n", - " \"\"\"\n", - " Search the vector database for semantically similar content.\n", - " \n", - " Args:\n", - " query: The search query text\n", - " limit: Maximum number of results to return\n", - " \n", - " Returns:\n", - " Search results as a formatted string\n", - " \"\"\"\n", - " # TODO: Implement actual vector search with embeddings\n", - " # For now, return a placeholder\n", - " return f\"Vector search for '{query}' (limit: {limit}) - Not yet implemented\"\n", - "\n", - "@agent.tool\n", - "async def query_graph(ctx: RunContext[DatabaseContext], gremlin_query: str) -> str:\n", - " \"\"\"\n", - " Execute a Gremlin query on the graph database.\n", - " \n", - " Args:\n", - " gremlin_query: The Gremlin query to execute\n", - " \n", - " Returns:\n", - " Query results as a formatted string\n", - " \"\"\"\n", - " try:\n", - " result = ctx.deps.janus_client.submit(gremlin_query).all().result()\n", - " return f\"Graph query results: {result}\"\n", - " except Exception as e:\n", - " return f\"Error executing graph query: {str(e)}\"\n", - "\n", - "@agent.tool\n", - "async def add_to_memory(ctx: RunContext[DatabaseContext], content: str, entity_type: str) -> str:\n", - " \"\"\"\n", - " Add new content to both vector and graph databases.\n", - " \n", - " Args:\n", - " content: The content to store\n", - " entity_type: The type of entity (e.g., 'job', 'company', 'person')\n", - " \n", - " Returns:\n", - " Confirmation message\n", - " \"\"\"\n", - " # TODO: Implement actual storage\n", - " # 1. Generate embedding for content\n", - " # 2. Store in Qdrant\n", - " # 3. Create node in JanusGraph\n", - " # 4. Link vector ID to graph node\n", - " return f\"Added {entity_type} to memory: '{content}' - Not yet implemented\"\n", - "\n", - "print(\"Agent tools registered:\")\n", - "print(\"- search_vector_memory: Search for semantically similar content\")\n", - "print(\"- query_graph: Execute Gremlin queries on the graph\")\n", - "print(\"- add_to_memory: Store new entities in the memory system\")" - ] + "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Test the Agent" - ] + ], + "id": "2bdc8896abdea4c3" }, { "cell_type": "code", - "execution_count": null, "metadata": {}, + "source": "# Initialize the memory system\nmemory = VectorGraphMemory(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_provider=embedding_provider,\n collection_name=\"job_search_memory\"\n)\n\n# Create database context\ndb_context = DatabaseContext(memory=memory)\n\n# Run a simple query\nresult = await agent.run(\n \"What tools do you have available?\",\n deps=db_context\n)\n\nprint(result.data)", + "id": "7d8b3e4682ffbf71", "outputs": [], - "source": [ - "# Create database context\n", - "db_context = DatabaseContext(\n", - " qdrant_client=qdrant,\n", - " janus_client=janus\n", - ")\n", - "\n", - "# Run a simple query\n", - "result = await agent.run(\n", - " \"What tools do you have available?\",\n", - " deps=db_context\n", - ")\n", - "\n", - "print(result.data)" - ] + "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Example: Job Search Use Case" - ] + ], + "id": "9c0e7713f32af3b0" }, { "cell_type": "code", - "execution_count": null, "metadata": {}, - "outputs": [], "source": [ "# Example: Add a job to memory\n", "result = await agent.run(\n", @@ -219,13 +129,14 @@ ")\n", "\n", "print(result.data)" - ] + ], + "id": "569ca3e38d410a94", + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "metadata": {}, - "outputs": [], "source": [ "# Example: Search for jobs\n", "result = await agent.run(\n", @@ -234,25 +145,30 @@ ")\n", "\n", "print(result.data)" - ] + ], + "id": "7d404fff0c65b3f6", + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Cleanup" - ] + ], + "id": "3da5d52101dd7132" }, { "cell_type": "code", - "execution_count": null, "metadata": {}, - "outputs": [], "source": [ "# Close connections\n", "janus.close()\n", "print(\"Database connections closed\")" - ] + ], + "id": "c28d83a09c52c4b4", + "outputs": [], + "execution_count": null } ], "metadata": { @@ -276,4 +192,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} +} \ No newline at end of file diff --git a/vector_graph_memory/__init__.py b/vector_graph_memory/__init__.py new file mode 100644 index 0000000..0d76881 --- /dev/null +++ b/vector_graph_memory/__init__.py @@ -0,0 +1,3 @@ +"""Vector Graph Memory - Hybrid vector-graph database backend for AI agents.""" + +__version__ = "0.1.0" diff --git a/vector_graph_memory/embeddings.py b/vector_graph_memory/embeddings.py new file mode 100644 index 0000000..9ee6fe4 --- /dev/null +++ b/vector_graph_memory/embeddings.py @@ -0,0 +1,126 @@ +"""Embedding provider interface and implementations.""" + +from abc import ABC, abstractmethod +from typing import List +import numpy as np + + +class EmbeddingProvider(ABC): + """Abstract interface for embedding providers.""" + + @abstractmethod + async def embed(self, text: str) -> List[float]: + """Generate embedding for a single text string. + + Args: + text: The text to embed + + Returns: + Embedding vector as a list of floats + """ + pass + + @abstractmethod + async def embed_batch(self, texts: List[str]) -> List[List[float]]: + """Generate embeddings for multiple text strings. + + Args: + texts: List of texts to embed + + Returns: + List of embedding vectors + """ + pass + + @property + @abstractmethod + def dimension(self) -> int: + """Return the dimensionality of the embeddings.""" + pass + + +class OpenAIEmbeddingProvider(EmbeddingProvider): + """OpenAI embedding provider using their API.""" + + def __init__(self, model: str = "text-embedding-3-small", api_key: str | None = None): + """Initialize OpenAI embedding provider. + + Args: + model: OpenAI embedding model name + api_key: OpenAI API key (if None, uses OPENAI_API_KEY env var) + """ + try: + from openai import AsyncOpenAI + except ImportError: + raise ImportError("openai package required for OpenAIEmbeddingProvider") + + self.model = model + self.client = AsyncOpenAI(api_key=api_key) + self._dimension = self._get_dimension(model) + + @staticmethod + def _get_dimension(model: str) -> int: + """Get embedding dimension for a given model.""" + dimensions = { + "text-embedding-3-small": 1536, + "text-embedding-3-large": 3072, + "text-embedding-ada-002": 1536, + } + return dimensions.get(model, 1536) + + async def embed(self, text: str) -> List[float]: + """Generate embedding for a single text string.""" + response = await self.client.embeddings.create( + model=self.model, + input=text + ) + return response.data[0].embedding + + async def embed_batch(self, texts: List[str]) -> List[List[float]]: + """Generate embeddings for multiple text strings.""" + response = await self.client.embeddings.create( + model=self.model, + input=texts + ) + return [item.embedding for item in response.data] + + @property + def dimension(self) -> int: + """Return the dimensionality of the embeddings.""" + return self._dimension + + +class SentenceTransformerEmbeddingProvider(EmbeddingProvider): + """Sentence Transformers embedding provider using local models.""" + + def __init__(self, model_name: str = "all-MiniLM-L6-v2"): + """Initialize Sentence Transformer embedding provider. + + Args: + model_name: HuggingFace model name for sentence transformers + """ + try: + from sentence_transformers import SentenceTransformer + except ImportError: + raise ImportError( + "sentence-transformers package required for SentenceTransformerEmbeddingProvider" + ) + + self.model = SentenceTransformer(model_name) + self._dimension = self.model.get_sentence_embedding_dimension() + + async def embed(self, text: str) -> List[float]: + """Generate embedding for a single text string.""" + # sentence-transformers is synchronous, but we wrap in async + embedding = self.model.encode(text, convert_to_numpy=True) + return embedding.tolist() + + async def embed_batch(self, texts: List[str]) -> List[List[float]]: + """Generate embeddings for multiple text strings.""" + embeddings = self.model.encode(texts, convert_to_numpy=True) + return embeddings.tolist() + + @property + def dimension(self) -> int: + """Return the dimensionality of the embeddings.""" + return self._dimension diff --git a/vector_graph_memory/memory.py b/vector_graph_memory/memory.py new file mode 100644 index 0000000..2ebfc8b --- /dev/null +++ b/vector_graph_memory/memory.py @@ -0,0 +1,199 @@ +"""Core memory system integrating vector and graph databases.""" + +import uuid +from typing import List, Dict, Any, Optional +from qdrant_client import QdrantClient +from qdrant_client.models import Distance, VectorParams, PointStruct, ScoredPoint +from gremlin_python.driver import client as gremlin_client + +from .embeddings import EmbeddingProvider + + +class VectorGraphMemory: + """Hybrid vector-graph memory system.""" + + def __init__( + self, + qdrant_client: QdrantClient, + janus_client: gremlin_client.Client, + embedding_provider: EmbeddingProvider, + collection_name: str = "memory", + ): + """Initialize vector-graph memory system. + + Args: + qdrant_client: Qdrant client instance + janus_client: JanusGraph Gremlin client instance + embedding_provider: Embedding provider for vectorization + collection_name: Name of the Qdrant collection + """ + self.qdrant = qdrant_client + self.janus = janus_client + self.embeddings = embedding_provider + self.collection_name = collection_name + + self._ensure_collection() + + def _ensure_collection(self) -> None: + """Ensure the Qdrant collection exists.""" + collections = self.qdrant.get_collections().collections + if not any(c.name == self.collection_name for c in collections): + self.qdrant.create_collection( + collection_name=self.collection_name, + vectors_config=VectorParams( + size=self.embeddings.dimension, + distance=Distance.COSINE + ) + ) + + async def add_entity( + self, + content: str, + entity_type: str, + metadata: Optional[Dict[str, Any]] = None + ) -> str: + """Add an entity to both vector and graph databases. + + Args: + content: The content/description of the entity + entity_type: Type of entity (e.g., 'job', 'company', 'person') + metadata: Additional metadata for the entity + + Returns: + Entity ID (UUID) + """ + entity_id = str(uuid.uuid4()) + metadata = metadata or {} + + # Generate embedding + embedding = await self.embeddings.embed(content) + + # Store in Qdrant + self.qdrant.upsert( + collection_name=self.collection_name, + points=[ + PointStruct( + id=entity_id, + vector=embedding, + payload={ + "content": content, + "type": entity_type, + **metadata + } + ) + ] + ) + + # Store in JanusGraph + gremlin_query = ( + f"g.addV('{entity_type}')" + f".property('id', '{entity_id}')" + f".property('content', content)" + ) + + # Add metadata properties + for key, value in metadata.items(): + # Escape single quotes in values + escaped_value = str(value).replace("'", "\\'") + gremlin_query += f".property('{key}', '{escaped_value}')" + + self.janus.submit(gremlin_query).all().result() + + return entity_id + + async def search_similar( + self, + query: str, + limit: int = 5, + entity_type: Optional[str] = None + ) -> List[Dict[str, Any]]: + """Search for semantically similar entities. + + Args: + query: Search query text + limit: Maximum number of results + entity_type: Optional filter by entity type + + Returns: + List of matching entities with scores + """ + # Generate query embedding + query_embedding = await self.embeddings.embed(query) + + # Search in Qdrant + search_kwargs = { + "collection_name": self.collection_name, + "query_vector": query_embedding, + "limit": limit + } + + if entity_type: + search_kwargs["query_filter"] = { + "must": [{"key": "type", "match": {"value": entity_type}}] + } + + results = self.qdrant.search(**search_kwargs) + + # Format results + return [ + { + "id": result.id, + "score": result.score, + "content": result.payload.get("content"), + "type": result.payload.get("type"), + "metadata": { + k: v for k, v in result.payload.items() + if k not in ["content", "type"] + } + } + for result in results + ] + + async def add_relationship( + self, + from_id: str, + to_id: str, + relationship_type: str, + properties: Optional[Dict[str, Any]] = None + ) -> None: + """Add a relationship between two entities in the graph. + + Args: + from_id: Source entity ID + to_id: Target entity ID + relationship_type: Type of relationship (e.g., 'works_at', 'applied_to') + properties: Optional properties for the relationship + """ + properties = properties or {} + + # Build Gremlin query to add edge + gremlin_query = ( + f"g.V().has('id', '{from_id}')" + f".addE('{relationship_type}')" + f".to(g.V().has('id', '{to_id}'))" + ) + + # Add properties to edge + for key, value in properties.items(): + escaped_value = str(value).replace("'", "\\'") + gremlin_query += f".property('{key}', '{escaped_value}')" + + self.janus.submit(gremlin_query).all().result() + + async def query_graph_traversal( + self, + start_id: str, + gremlin_steps: str + ) -> List[Dict[str, Any]]: + """Execute a graph traversal starting from a specific entity. + + Args: + start_id: Starting entity ID + gremlin_steps: Gremlin traversal steps to execute + + Returns: + List of results from the traversal + """ + gremlin_query = f"g.V().has('id', '{start_id}').{gremlin_steps}" + results = self.janus.submit(gremlin_query).all().result() + return results From 242bd14b32d58d51693dfa2228efb66f236e8da6 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 16:21:38 -0700 Subject: [PATCH 10/21] refactor: use PydanticAI native embedding models Remove custom embedding provider abstraction in favor of PydanticAI's built-in EmbeddingModel interface: - Delete vector_graph_memory/embeddings.py (custom providers) - Update VectorGraphMemory to accept PydanticAI EmbeddingModel - Simplify model configuration to use PydanticAI's infer_model() Update playground notebook: - Use pydantic_ai.models.infer_model() for embedding configuration - Single EMBEDDING_MODEL env var (e.g., 'openai:text-embedding-3-small') - Remove provider-specific conditional logic Add MODEL_SELECTION.md guide: - How to choose LLM models (OpenAI, Anthropic, Google, Groq, Ollama) - How to choose embedding models (OpenAI, Google, Cohere, VoyageAI) - Configuration examples and best practices - API key setup instructions Benefits: - Leverages PydanticAI's native model support - Single configuration pattern for both LLM and embeddings - Automatic dimension detection - Support for all PydanticAI-compatible providers --- .env.example | 21 ----- .gitignore | 4 + MODEL_SELECTION.md | 127 ++++++++++++++++++++++++++++++ playground.ipynb | 6 +- pyproject.toml | 2 + vector_graph_memory/embeddings.py | 126 ----------------------------- vector_graph_memory/memory.py | 28 ++++--- 7 files changed, 153 insertions(+), 161 deletions(-) delete mode 100644 .env.example create mode 100644 MODEL_SELECTION.md delete mode 100644 vector_graph_memory/embeddings.py diff --git a/.env.example b/.env.example deleted file mode 100644 index e65298b..0000000 --- a/.env.example +++ /dev/null @@ -1,21 +0,0 @@ -# Docker Compose Configuration -# These variables control the Docker containers - -# Qdrant Vector Database -QDRANT_VERSION=latest -QDRANT_CONTAINER_NAME=vector-graph-qdrant -QDRANT_HTTP_PORT=6333 -QDRANT_GRPC_PORT=6334 - -# JanusGraph Graph Database -JANUSGRAPH_VERSION=latest -JANUSGRAPH_CONTAINER_NAME=vector-graph-janusgraph -JANUSGRAPH_PORT=8182 - -# Application Configuration -# These variables are used by the Python application to connect to databases - -QDRANT_HOST=localhost -QDRANT_PORT=6333 - -JANUSGRAPH_HOST=localhost diff --git a/.gitignore b/.gitignore index e8fece7..b5958de 100644 --- a/.gitignore +++ b/.gitignore @@ -205,3 +205,7 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + + +## Custom Ignores +uv.lock diff --git a/MODEL_SELECTION.md b/MODEL_SELECTION.md new file mode 100644 index 0000000..dfe87f0 --- /dev/null +++ b/MODEL_SELECTION.md @@ -0,0 +1,127 @@ +# Model Selection Guide + +This guide explains how to configure LLM and embedding models for Vector Graph Memory. + +## Quick Start + +Configure models via environment variables in `.env`: + +```bash +# LLM Model for the agent +LLM_MODEL=openai:gpt-4 + +# Embedding Model for vectorization +EMBEDDING_MODEL=openai:text-embedding-3-small +``` + +## LLM Models + +Vector Graph Memory uses PydanticAI's agent system, which supports multiple LLM providers. + +### Supported Providers + +| Provider | Format | Example | +|----------|--------|---------| +| OpenAI | `openai:` | `openai:gpt-4` | +| Anthropic | `anthropic:` | `anthropic:claude-3-5-sonnet-20241022` | +| Google Gemini | `google-gla:` | `google-gla:gemini-1.5-pro` | +| Google Vertex | `google-vertex:` | `google-vertex:gemini-1.5-pro` | +| Groq | `groq:` | `groq:llama-3.1-70b-versatile` | +| Ollama | `ollama:` | `ollama:llama3` | + +### Examples + +```bash +# Use OpenAI GPT-4 +LLM_MODEL=openai:gpt-4 + +# Use Anthropic Claude +LLM_MODEL=anthropic:claude-3-5-sonnet-20241022 + +# Use local Ollama model +LLM_MODEL=ollama:llama3 +``` + +## Embedding Models + +Embedding models are used to convert text into vector representations for semantic search. + +### Supported Providers + +| Provider | Format | Example | +|----------|--------|---------| +| OpenAI | `openai:` | `openai:text-embedding-3-small` | +| Google Gemini | `google-gla:` | `google-gla:text-embedding-004` | +| Google Vertex | `google-vertex:` | `google-vertex:text-embedding-004` | +| Cohere | `cohere:` | `cohere:embed-english-v3.0` | +| VoyageAI | `voyageai:` | `voyageai:voyage-2` | + +### Choosing an Embedding Model + +**For general use:** +- `openai:text-embedding-3-small` - Good balance of cost and performance (1536 dimensions) +- `openai:text-embedding-3-large` - Higher quality (3072 dimensions) + +**For multilingual support:** +- `cohere:embed-multilingual-v3.0` - Supports 100+ languages + +**For domain-specific tasks:** +- `voyageai:voyage-code-2` - Optimized for code +- `voyageai:voyage-finance-2` - Optimized for financial documents + +### Examples + +```bash +# Use OpenAI small embedding model (recommended for most use cases) +EMBEDDING_MODEL=openai:text-embedding-3-small + +# Use Google's embedding model +EMBEDDING_MODEL=google-gla:text-embedding-004 + +# Use Cohere for multilingual support +EMBEDDING_MODEL=cohere:embed-multilingual-v3.0 +``` + +## API Keys + +Set API keys as environment variables: + +```bash +# OpenAI +OPENAI_API_KEY=your-api-key + +# Anthropic +ANTHROPIC_API_KEY=your-api-key + +# Google (Gemini) +GOOGLE_API_KEY=your-api-key + +# Cohere +COHERE_API_KEY=your-api-key + +# VoyageAI +VOYAGEAI_API_KEY=your-api-key +``` + +## Complete Example + +```bash +# .env file example +LLM_MODEL=anthropic:claude-3-5-sonnet-20241022 +EMBEDDING_MODEL=openai:text-embedding-3-small + +ANTHROPIC_API_KEY=your-anthropic-key +OPENAI_API_KEY=your-openai-key +``` + +## Notes + +- **LLM and embedding models can use different providers** - For example, you can use Claude for the agent and OpenAI for embeddings +- **Model strings follow PydanticAI's format** - See [PydanticAI documentation](https://ai.pydantic.dev/models/overview/) for details +- **Embedding dimensions are detected automatically** - The system will create Qdrant collections with the correct vector size +- **Local models via Ollama** - You can use local models for both LLM and embeddings to avoid API costs + +## Further Reading + +- [PydanticAI Models Documentation](https://ai.pydantic.dev/models/overview/) +- [PydanticAI Embeddings Documentation](https://ai.pydantic.dev/embeddings/) diff --git a/playground.ipynb b/playground.ipynb index 907c50b..deb7514 100644 --- a/playground.ipynb +++ b/playground.ipynb @@ -26,7 +26,7 @@ "start_time": "2026-03-21T23:12:23.803285021Z" } }, - "source": "import os\nfrom dotenv import load_dotenv\nfrom qdrant_client import QdrantClient\nfrom gremlin_python.driver import client as gremlin_client\nfrom pydantic_ai import Agent, RunContext\nfrom pydantic import BaseModel\n\nfrom vector_graph_memory.embeddings import OpenAIEmbeddingProvider, SentenceTransformerEmbeddingProvider\nfrom vector_graph_memory.memory import VectorGraphMemory\n\n# Load environment variables\nload_dotenv()", + "source": "import os\nfrom dotenv import load_dotenv\nfrom qdrant_client import QdrantClient\nfrom gremlin_python.driver import client as gremlin_client\nfrom pydantic_ai import Agent, RunContext\nfrom pydantic_ai.models import infer_model\nfrom pydantic import BaseModel\n\nfrom vector_graph_memory.memory import VectorGraphMemory\n\n# Load environment variables\nload_dotenv()", "id": "dd65b95cb0cd2644", "outputs": [], "execution_count": null @@ -57,7 +57,7 @@ { "cell_type": "code", "metadata": {}, - "source": "# JanusGraph connection\njanusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\njanusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n\njanus = gremlin_client.Client(\n f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n 'g'\n)\nprint(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")\n\n# Initialize embedding provider (configurable)\nembedding_provider_type = os.getenv('EMBEDDING_PROVIDER', 'openai')\n\nif embedding_provider_type == 'openai':\n embedding_provider = OpenAIEmbeddingProvider(\n model=os.getenv('OPENAI_EMBEDDING_MODEL', 'text-embedding-3-small')\n )\nelif embedding_provider_type == 'sentence-transformers':\n embedding_provider = SentenceTransformerEmbeddingProvider(\n model_name=os.getenv('SENTENCE_TRANSFORMER_MODEL', 'all-MiniLM-L6-v2')\n )\nelse:\n raise ValueError(f\"Unknown embedding provider: {embedding_provider_type}\")\n\nprint(f\"Initialized {embedding_provider_type} embedding provider (dim={embedding_provider.dimension})\")", + "source": "# JanusGraph connection\njanusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\njanusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n\njanus = gremlin_client.Client(\n f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n 'g'\n)\nprint(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")\n\n# Initialize embedding model using PydanticAI\n# Supported formats: 'openai:text-embedding-3-small', 'google-gla:text-embedding-004', etc.\nembedding_model_name = os.getenv('EMBEDDING_MODEL', 'openai:text-embedding-3-small')\nembedding_model = infer_model(embedding_model_name)\n\nprint(f\"Initialized embedding model: {embedding_model_name}\")", "id": "173bac3780209c1d", "outputs": [], "execution_count": null @@ -105,7 +105,7 @@ { "cell_type": "code", "metadata": {}, - "source": "# Initialize the memory system\nmemory = VectorGraphMemory(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_provider=embedding_provider,\n collection_name=\"job_search_memory\"\n)\n\n# Create database context\ndb_context = DatabaseContext(memory=memory)\n\n# Run a simple query\nresult = await agent.run(\n \"What tools do you have available?\",\n deps=db_context\n)\n\nprint(result.data)", + "source": "# Initialize the memory system\nmemory = VectorGraphMemory(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_model=embedding_model,\n collection_name=\"job_search_memory\"\n)\n\n# Create database context\ndb_context = DatabaseContext(memory=memory)\n\n# Run a simple query\nresult = await agent.run(\n \"What tools do you have available?\",\n deps=db_context\n)\n\nprint(result.data)", "id": "7d8b3e4682ffbf71", "outputs": [], "execution_count": null diff --git a/pyproject.toml b/pyproject.toml index 67de69b..a848a89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,9 @@ version = "0.1.0" description = "Add your description here" requires-python = ">=3.11" dependencies = [ + "gremlinpython>=3.8.0", "pydantic-ai>=1.70.0", + "qdrant-client>=1.17.1", ] [dependency-groups] diff --git a/vector_graph_memory/embeddings.py b/vector_graph_memory/embeddings.py deleted file mode 100644 index 9ee6fe4..0000000 --- a/vector_graph_memory/embeddings.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Embedding provider interface and implementations.""" - -from abc import ABC, abstractmethod -from typing import List -import numpy as np - - -class EmbeddingProvider(ABC): - """Abstract interface for embedding providers.""" - - @abstractmethod - async def embed(self, text: str) -> List[float]: - """Generate embedding for a single text string. - - Args: - text: The text to embed - - Returns: - Embedding vector as a list of floats - """ - pass - - @abstractmethod - async def embed_batch(self, texts: List[str]) -> List[List[float]]: - """Generate embeddings for multiple text strings. - - Args: - texts: List of texts to embed - - Returns: - List of embedding vectors - """ - pass - - @property - @abstractmethod - def dimension(self) -> int: - """Return the dimensionality of the embeddings.""" - pass - - -class OpenAIEmbeddingProvider(EmbeddingProvider): - """OpenAI embedding provider using their API.""" - - def __init__(self, model: str = "text-embedding-3-small", api_key: str | None = None): - """Initialize OpenAI embedding provider. - - Args: - model: OpenAI embedding model name - api_key: OpenAI API key (if None, uses OPENAI_API_KEY env var) - """ - try: - from openai import AsyncOpenAI - except ImportError: - raise ImportError("openai package required for OpenAIEmbeddingProvider") - - self.model = model - self.client = AsyncOpenAI(api_key=api_key) - self._dimension = self._get_dimension(model) - - @staticmethod - def _get_dimension(model: str) -> int: - """Get embedding dimension for a given model.""" - dimensions = { - "text-embedding-3-small": 1536, - "text-embedding-3-large": 3072, - "text-embedding-ada-002": 1536, - } - return dimensions.get(model, 1536) - - async def embed(self, text: str) -> List[float]: - """Generate embedding for a single text string.""" - response = await self.client.embeddings.create( - model=self.model, - input=text - ) - return response.data[0].embedding - - async def embed_batch(self, texts: List[str]) -> List[List[float]]: - """Generate embeddings for multiple text strings.""" - response = await self.client.embeddings.create( - model=self.model, - input=texts - ) - return [item.embedding for item in response.data] - - @property - def dimension(self) -> int: - """Return the dimensionality of the embeddings.""" - return self._dimension - - -class SentenceTransformerEmbeddingProvider(EmbeddingProvider): - """Sentence Transformers embedding provider using local models.""" - - def __init__(self, model_name: str = "all-MiniLM-L6-v2"): - """Initialize Sentence Transformer embedding provider. - - Args: - model_name: HuggingFace model name for sentence transformers - """ - try: - from sentence_transformers import SentenceTransformer - except ImportError: - raise ImportError( - "sentence-transformers package required for SentenceTransformerEmbeddingProvider" - ) - - self.model = SentenceTransformer(model_name) - self._dimension = self.model.get_sentence_embedding_dimension() - - async def embed(self, text: str) -> List[float]: - """Generate embedding for a single text string.""" - # sentence-transformers is synchronous, but we wrap in async - embedding = self.model.encode(text, convert_to_numpy=True) - return embedding.tolist() - - async def embed_batch(self, texts: List[str]) -> List[List[float]]: - """Generate embeddings for multiple text strings.""" - embeddings = self.model.encode(texts, convert_to_numpy=True) - return embeddings.tolist() - - @property - def dimension(self) -> int: - """Return the dimensionality of the embeddings.""" - return self._dimension diff --git a/vector_graph_memory/memory.py b/vector_graph_memory/memory.py index 2ebfc8b..e1270e8 100644 --- a/vector_graph_memory/memory.py +++ b/vector_graph_memory/memory.py @@ -3,10 +3,9 @@ import uuid from typing import List, Dict, Any, Optional from qdrant_client import QdrantClient -from qdrant_client.models import Distance, VectorParams, PointStruct, ScoredPoint +from qdrant_client.models import Distance, VectorParams, PointStruct from gremlin_python.driver import client as gremlin_client - -from .embeddings import EmbeddingProvider +from pydantic_ai.models import EmbeddingModel class VectorGraphMemory: @@ -16,7 +15,7 @@ def __init__( self, qdrant_client: QdrantClient, janus_client: gremlin_client.Client, - embedding_provider: EmbeddingProvider, + embedding_model: EmbeddingModel, collection_name: str = "memory", ): """Initialize vector-graph memory system. @@ -24,12 +23,12 @@ def __init__( Args: qdrant_client: Qdrant client instance janus_client: JanusGraph Gremlin client instance - embedding_provider: Embedding provider for vectorization + embedding_model: PydanticAI EmbeddingModel instance collection_name: Name of the Qdrant collection """ self.qdrant = qdrant_client self.janus = janus_client - self.embeddings = embedding_provider + self.embedding_model = embedding_model self.collection_name = collection_name self._ensure_collection() @@ -38,10 +37,15 @@ def _ensure_collection(self) -> None: """Ensure the Qdrant collection exists.""" collections = self.qdrant.get_collections().collections if not any(c.name == self.collection_name for c in collections): + # Get embedding dimension from a test embedding + import asyncio + test_result = asyncio.run(self.embedding_model.embed(["test"])) + embedding_dim = len(test_result.embeddings[0]) + self.qdrant.create_collection( collection_name=self.collection_name, vectors_config=VectorParams( - size=self.embeddings.dimension, + size=embedding_dim, distance=Distance.COSINE ) ) @@ -65,8 +69,9 @@ async def add_entity( entity_id = str(uuid.uuid4()) metadata = metadata or {} - # Generate embedding - embedding = await self.embeddings.embed(content) + # Generate embedding using PydanticAI + embedding_result = await self.embedding_model.embed([content]) + embedding = embedding_result.embeddings[0] # Store in Qdrant self.qdrant.upsert( @@ -117,8 +122,9 @@ async def search_similar( Returns: List of matching entities with scores """ - # Generate query embedding - query_embedding = await self.embeddings.embed(query) + # Generate query embedding using PydanticAI + embedding_result = await self.embedding_model.embed([query]) + query_embedding = embedding_result.embeddings[0] # Search in Qdrant search_kwargs = { From 934d1db45317738dc134038f1bede7a7b113ba97 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 16:23:39 -0700 Subject: [PATCH 11/21] refactor: make models API parameters not config Change model selection from environment variables to API parameters: Playground notebook: - Pass 'openai:gpt-4' directly to Agent() - Pass infer_model('openai:text-embedding-3-small') to VectorGraphMemory() - Remove env var lookups for LLM_MODEL and EMBEDDING_MODEL - Add inline comments showing examples of other providers MODEL_SELECTION.md: - Rewrite to show code examples instead of env var configuration - Emphasize "models are API parameters, not configuration" - Show how to mix providers (e.g., Claude LLM + OpenAI embeddings) - Add "Key Design Principles" section This is the correct design pattern - users pass model strings directly to the library API rather than configuring via environment variables. --- MODEL_SELECTION.md | 124 +++++++++++++++++++++++++++++---------------- playground.ipynb | 6 +-- 2 files changed, 82 insertions(+), 48 deletions(-) diff --git a/MODEL_SELECTION.md b/MODEL_SELECTION.md index dfe87f0..9f5e533 100644 --- a/MODEL_SELECTION.md +++ b/MODEL_SELECTION.md @@ -1,22 +1,29 @@ # Model Selection Guide -This guide explains how to configure LLM and embedding models for Vector Graph Memory. +This guide explains how to choose LLM and embedding models for Vector Graph Memory. ## Quick Start -Configure models via environment variables in `.env`: - -```bash -# LLM Model for the agent -LLM_MODEL=openai:gpt-4 - -# Embedding Model for vectorization -EMBEDDING_MODEL=openai:text-embedding-3-small +```python +from pydantic_ai import Agent +from pydantic_ai.models import infer_model +from vector_graph_memory import VectorGraphMemory + +# Choose your LLM for the agent +agent = Agent('openai:gpt-4', deps_type=DatabaseContext) + +# Choose your embedding model for memory +embedding_model = infer_model('openai:text-embedding-3-small') +memory = VectorGraphMemory( + qdrant_client=qdrant, + janus_client=janus, + embedding_model=embedding_model +) ``` ## LLM Models -Vector Graph Memory uses PydanticAI's agent system, which supports multiple LLM providers. +The LLM model is passed directly to `Agent()` from PydanticAI. This is the model that powers your AI agent's reasoning and tool use. ### Supported Providers @@ -31,20 +38,20 @@ Vector Graph Memory uses PydanticAI's agent system, which supports multiple LLM ### Examples -```bash +```python # Use OpenAI GPT-4 -LLM_MODEL=openai:gpt-4 +agent = Agent('openai:gpt-4') # Use Anthropic Claude -LLM_MODEL=anthropic:claude-3-5-sonnet-20241022 +agent = Agent('anthropic:claude-3-5-sonnet-20241022') # Use local Ollama model -LLM_MODEL=ollama:llama3 +agent = Agent('ollama:llama3') ``` ## Embedding Models -Embedding models are used to convert text into vector representations for semantic search. +The embedding model is passed to `VectorGraphMemory()` to convert text into vectors for semantic search. ### Supported Providers @@ -71,15 +78,55 @@ Embedding models are used to convert text into vector representations for semant ### Examples -```bash +```python +from pydantic_ai.models import infer_model + # Use OpenAI small embedding model (recommended for most use cases) -EMBEDDING_MODEL=openai:text-embedding-3-small +embedding_model = infer_model('openai:text-embedding-3-small') # Use Google's embedding model -EMBEDDING_MODEL=google-gla:text-embedding-004 +embedding_model = infer_model('google-gla:text-embedding-004') # Use Cohere for multilingual support -EMBEDDING_MODEL=cohere:embed-multilingual-v3.0 +embedding_model = infer_model('cohere:embed-multilingual-v3.0') + +# Pass to VectorGraphMemory +memory = VectorGraphMemory( + qdrant_client=qdrant, + janus_client=janus, + embedding_model=embedding_model +) +``` + +## Complete Example + +```python +from pydantic_ai import Agent, RunContext +from pydantic_ai.models import infer_model +from qdrant_client import QdrantClient +from gremlin_python.driver import client as gremlin_client +from vector_graph_memory import VectorGraphMemory + +# Connect to databases +qdrant = QdrantClient(host='localhost', port=6333) +janus = gremlin_client.Client('ws://localhost:8182/gremlin', 'g') + +# Choose embedding model +embedding_model = infer_model('openai:text-embedding-3-small') + +# Initialize memory system +memory = VectorGraphMemory( + qdrant_client=qdrant, + janus_client=janus, + embedding_model=embedding_model, + collection_name="my_memory" +) + +# Create agent with chosen LLM +agent = Agent('anthropic:claude-3-5-sonnet-20241022') + +# Use different providers for LLM vs embeddings! +# For example: Claude for reasoning + OpenAI for embeddings ``` ## API Keys @@ -87,39 +134,26 @@ EMBEDDING_MODEL=cohere:embed-multilingual-v3.0 Set API keys as environment variables: ```bash -# OpenAI -OPENAI_API_KEY=your-api-key - -# Anthropic -ANTHROPIC_API_KEY=your-api-key - -# Google (Gemini) -GOOGLE_API_KEY=your-api-key - -# Cohere -COHERE_API_KEY=your-api-key - -# VoyageAI -VOYAGEAI_API_KEY=your-api-key +export OPENAI_API_KEY=your-api-key +export ANTHROPIC_API_KEY=your-api-key +export GOOGLE_API_KEY=your-api-key +export COHERE_API_KEY=your-api-key +export VOYAGEAI_API_KEY=your-api-key ``` -## Complete Example +Or in a `.env` file: ```bash -# .env file example -LLM_MODEL=anthropic:claude-3-5-sonnet-20241022 -EMBEDDING_MODEL=openai:text-embedding-3-small - -ANTHROPIC_API_KEY=your-anthropic-key -OPENAI_API_KEY=your-openai-key +OPENAI_API_KEY=your-api-key +ANTHROPIC_API_KEY=your-api-key ``` -## Notes +## Key Design Principles -- **LLM and embedding models can use different providers** - For example, you can use Claude for the agent and OpenAI for embeddings -- **Model strings follow PydanticAI's format** - See [PydanticAI documentation](https://ai.pydantic.dev/models/overview/) for details -- **Embedding dimensions are detected automatically** - The system will create Qdrant collections with the correct vector size -- **Local models via Ollama** - You can use local models for both LLM and embeddings to avoid API costs +1. **Models are API parameters, not configuration** - Pass model strings directly to `Agent()` and `infer_model()` +2. **Mix and match providers** - Use different providers for LLM vs embeddings (e.g., Claude + OpenAI) +3. **Drop-in replacement** - Change model strings to switch providers without code changes +4. **PydanticAI native** - Leverages PydanticAI's built-in model support ## Further Reading diff --git a/playground.ipynb b/playground.ipynb index deb7514..83245e8 100644 --- a/playground.ipynb +++ b/playground.ipynb @@ -57,7 +57,7 @@ { "cell_type": "code", "metadata": {}, - "source": "# JanusGraph connection\njanusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\njanusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n\njanus = gremlin_client.Client(\n f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n 'g'\n)\nprint(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")\n\n# Initialize embedding model using PydanticAI\n# Supported formats: 'openai:text-embedding-3-small', 'google-gla:text-embedding-004', etc.\nembedding_model_name = os.getenv('EMBEDDING_MODEL', 'openai:text-embedding-3-small')\nembedding_model = infer_model(embedding_model_name)\n\nprint(f\"Initialized embedding model: {embedding_model_name}\")", + "source": "# JanusGraph connection\njanusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\njanusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n\njanus = gremlin_client.Client(\n f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n 'g'\n)\nprint(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")", "id": "173bac3780209c1d", "outputs": [], "execution_count": null @@ -89,7 +89,7 @@ { "cell_type": "code", "metadata": {}, - "source": "# Create the agent with configurable LLM model\nllm_model = os.getenv('LLM_MODEL', 'openai:gpt-4')\n\nagent = Agent(\n llm_model,\n deps_type=DatabaseContext,\n system_prompt=\"You are a helpful assistant with access to a vector-graph memory system.\"\n)\n\n@agent.tool\nasync def search_vector_memory(ctx: RunContext[DatabaseContext], query: str, limit: int = 5) -> str:\n \"\"\"\n Search the vector database for semantically similar content.\n \n Args:\n query: The search query text\n limit: Maximum number of results to return\n \n Returns:\n Search results as a formatted string\n \"\"\"\n results = await ctx.deps.memory.search_similar(query, limit=limit)\n \n if not results:\n return f\"No results found for query: '{query}'\"\n \n formatted_results = []\n for i, result in enumerate(results, 1):\n formatted_results.append(\n f\"{i}. [{result['type']}] {result['content']} (score: {result['score']:.3f})\"\n )\n \n return \"\\\\n\".join(formatted_results)\n\n@agent.tool\nasync def query_graph(ctx: RunContext[DatabaseContext], start_id: str, traversal_steps: str) -> str:\n \"\"\"\n Execute a graph traversal starting from a specific entity.\n \n Args:\n start_id: Starting entity ID\n traversal_steps: Gremlin traversal steps (e.g., \"out('works_at').values('content')\")\n \n Returns:\n Query results as a formatted string\n \"\"\"\n try:\n results = await ctx.deps.memory.query_graph_traversal(start_id, traversal_steps)\n return f\"Graph traversal results: {results}\"\n except Exception as e:\n return f\"Error executing graph traversal: {str(e)}\"\n\n@agent.tool\nasync def add_to_memory(ctx: RunContext[DatabaseContext], content: str, entity_type: str, metadata: dict | None = None) -> str:\n \"\"\"\n Add new content to both vector and graph databases.\n \n Args:\n content: The content to store\n entity_type: The type of entity (e.g., 'job', 'company', 'person')\n metadata: Optional additional metadata\n \n Returns:\n Confirmation message with entity ID\n \"\"\"\n entity_id = await ctx.deps.memory.add_entity(content, entity_type, metadata)\n return f\"Added {entity_type} to memory with ID: {entity_id}\"\n\n@agent.tool\nasync def link_entities(ctx: RunContext[DatabaseContext], from_id: str, to_id: str, relationship: str) -> str:\n \"\"\"\n Create a relationship between two entities in the graph.\n \n Args:\n from_id: Source entity ID\n to_id: Target entity ID\n relationship: Relationship type (e.g., 'works_at', 'applied_to')\n \n Returns:\n Confirmation message\n \"\"\"\n try:\n await ctx.deps.memory.add_relationship(from_id, to_id, relationship)\n return f\"Created '{relationship}' relationship from {from_id} to {to_id}\"\n except Exception as e:\n return f\"Error creating relationship: {str(e)}\"\n\nprint(f\"Agent initialized with LLM: {llm_model}\")\nprint(\"Agent tools registered:\")\nprint(\"- search_vector_memory: Search for semantically similar content\")\nprint(\"- query_graph: Execute graph traversals from specific entities\")\nprint(\"- add_to_memory: Store new entities in the memory system\")\nprint(\"- link_entities: Create relationships between entities\")", + "source": "# Create the agent - choose your LLM model\n# Examples: 'openai:gpt-4', 'anthropic:claude-3-5-sonnet-20241022', 'ollama:llama3'\nagent = Agent(\n 'openai:gpt-4',\n deps_type=DatabaseContext,\n system_prompt=\"You are a helpful assistant with access to a vector-graph memory system.\"\n)\n\n@agent.tool\nasync def search_vector_memory(ctx: RunContext[DatabaseContext], query: str, limit: int = 5) -> str:\n \"\"\"\n Search the vector database for semantically similar content.\n \n Args:\n query: The search query text\n limit: Maximum number of results to return\n \n Returns:\n Search results as a formatted string\n \"\"\"\n results = await ctx.deps.memory.search_similar(query, limit=limit)\n \n if not results:\n return f\"No results found for query: '{query}'\"\n \n formatted_results = []\n for i, result in enumerate(results, 1):\n formatted_results.append(\n f\"{i}. [{result['type']}] {result['content']} (score: {result['score']:.3f})\"\n )\n \n return \"\\\\n\".join(formatted_results)\n\n@agent.tool\nasync def query_graph(ctx: RunContext[DatabaseContext], start_id: str, traversal_steps: str) -> str:\n \"\"\"\n Execute a graph traversal starting from a specific entity.\n \n Args:\n start_id: Starting entity ID\n traversal_steps: Gremlin traversal steps (e.g., \"out('works_at').values('content')\")\n \n Returns:\n Query results as a formatted string\n \"\"\"\n try:\n results = await ctx.deps.memory.query_graph_traversal(start_id, traversal_steps)\n return f\"Graph traversal results: {results}\"\n except Exception as e:\n return f\"Error executing graph traversal: {str(e)}\"\n\n@agent.tool\nasync def add_to_memory(ctx: RunContext[DatabaseContext], content: str, entity_type: str, metadata: dict | None = None) -> str:\n \"\"\"\n Add new content to both vector and graph databases.\n \n Args:\n content: The content to store\n entity_type: The type of entity (e.g., 'job', 'company', 'person')\n metadata: Optional additional metadata\n \n Returns:\n Confirmation message with entity ID\n \"\"\"\n entity_id = await ctx.deps.memory.add_entity(content, entity_type, metadata)\n return f\"Added {entity_type} to memory with ID: {entity_id}\"\n\n@agent.tool\nasync def link_entities(ctx: RunContext[DatabaseContext], from_id: str, to_id: str, relationship: str) -> str:\n \"\"\"\n Create a relationship between two entities in the graph.\n \n Args:\n from_id: Source entity ID\n to_id: Target entity ID\n relationship: Relationship type (e.g., 'works_at', 'applied_to')\n \n Returns:\n Confirmation message\n \"\"\"\n try:\n await ctx.deps.memory.add_relationship(from_id, to_id, relationship)\n return f\"Created '{relationship}' relationship from {from_id} to {to_id}\"\n except Exception as e:\n return f\"Error creating relationship: {str(e)}\"\n\nprint(\"Agent initialized with LLM: openai:gpt-4\")\nprint(\"Agent tools registered:\")\nprint(\"- search_vector_memory: Search for semantically similar content\")\nprint(\"- query_graph: Execute graph traversals from specific entities\")\nprint(\"- add_to_memory: Store new entities in the memory system\")\nprint(\"- link_entities: Create relationships between entities\")", "id": "832d50545225edcf", "outputs": [], "execution_count": null @@ -105,7 +105,7 @@ { "cell_type": "code", "metadata": {}, - "source": "# Initialize the memory system\nmemory = VectorGraphMemory(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_model=embedding_model,\n collection_name=\"job_search_memory\"\n)\n\n# Create database context\ndb_context = DatabaseContext(memory=memory)\n\n# Run a simple query\nresult = await agent.run(\n \"What tools do you have available?\",\n deps=db_context\n)\n\nprint(result.data)", + "source": "# Initialize the memory system - choose your embedding model\n# Examples: 'openai:text-embedding-3-small', 'google-gla:text-embedding-004', 'cohere:embed-english-v3.0'\nembedding_model = infer_model('openai:text-embedding-3-small')\n\nmemory = VectorGraphMemory(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_model=embedding_model,\n collection_name=\"job_search_memory\"\n)\n\n# Create database context\ndb_context = DatabaseContext(memory=memory)\n\nprint(\"Memory system initialized with embedding model: openai:text-embedding-3-small\")\n\n# Run a simple query\nresult = await agent.run(\n \"What tools do you have available?\",\n deps=db_context\n)\n\nprint(result.data)", "id": "7d8b3e4682ffbf71", "outputs": [], "execution_count": null From 5fb16e55130c904ffbbff68c582a782a0979a5d1 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:03 -0700 Subject: [PATCH 12/21] refactor: reorganize package structure and remove outdated docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Move from vector_graph_memory/ to src/vgm/ following src-layout - Delete MODEL_SELECTION.md (model config now in API parameters) - Delete old memory.py implementation - Prepare for new modular architecture 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- MODEL_SELECTION.md | 161 ------------------------- src/vgm/__init__.py | 54 +++++++++ vector_graph_memory/__init__.py | 3 - vector_graph_memory/memory.py | 205 -------------------------------- 4 files changed, 54 insertions(+), 369 deletions(-) delete mode 100644 MODEL_SELECTION.md create mode 100644 src/vgm/__init__.py delete mode 100644 vector_graph_memory/__init__.py delete mode 100644 vector_graph_memory/memory.py diff --git a/MODEL_SELECTION.md b/MODEL_SELECTION.md deleted file mode 100644 index 9f5e533..0000000 --- a/MODEL_SELECTION.md +++ /dev/null @@ -1,161 +0,0 @@ -# Model Selection Guide - -This guide explains how to choose LLM and embedding models for Vector Graph Memory. - -## Quick Start - -```python -from pydantic_ai import Agent -from pydantic_ai.models import infer_model -from vector_graph_memory import VectorGraphMemory - -# Choose your LLM for the agent -agent = Agent('openai:gpt-4', deps_type=DatabaseContext) - -# Choose your embedding model for memory -embedding_model = infer_model('openai:text-embedding-3-small') -memory = VectorGraphMemory( - qdrant_client=qdrant, - janus_client=janus, - embedding_model=embedding_model -) -``` - -## LLM Models - -The LLM model is passed directly to `Agent()` from PydanticAI. This is the model that powers your AI agent's reasoning and tool use. - -### Supported Providers - -| Provider | Format | Example | -|----------|--------|---------| -| OpenAI | `openai:` | `openai:gpt-4` | -| Anthropic | `anthropic:` | `anthropic:claude-3-5-sonnet-20241022` | -| Google Gemini | `google-gla:` | `google-gla:gemini-1.5-pro` | -| Google Vertex | `google-vertex:` | `google-vertex:gemini-1.5-pro` | -| Groq | `groq:` | `groq:llama-3.1-70b-versatile` | -| Ollama | `ollama:` | `ollama:llama3` | - -### Examples - -```python -# Use OpenAI GPT-4 -agent = Agent('openai:gpt-4') - -# Use Anthropic Claude -agent = Agent('anthropic:claude-3-5-sonnet-20241022') - -# Use local Ollama model -agent = Agent('ollama:llama3') -``` - -## Embedding Models - -The embedding model is passed to `VectorGraphMemory()` to convert text into vectors for semantic search. - -### Supported Providers - -| Provider | Format | Example | -|----------|--------|---------| -| OpenAI | `openai:` | `openai:text-embedding-3-small` | -| Google Gemini | `google-gla:` | `google-gla:text-embedding-004` | -| Google Vertex | `google-vertex:` | `google-vertex:text-embedding-004` | -| Cohere | `cohere:` | `cohere:embed-english-v3.0` | -| VoyageAI | `voyageai:` | `voyageai:voyage-2` | - -### Choosing an Embedding Model - -**For general use:** -- `openai:text-embedding-3-small` - Good balance of cost and performance (1536 dimensions) -- `openai:text-embedding-3-large` - Higher quality (3072 dimensions) - -**For multilingual support:** -- `cohere:embed-multilingual-v3.0` - Supports 100+ languages - -**For domain-specific tasks:** -- `voyageai:voyage-code-2` - Optimized for code -- `voyageai:voyage-finance-2` - Optimized for financial documents - -### Examples - -```python -from pydantic_ai.models import infer_model - -# Use OpenAI small embedding model (recommended for most use cases) -embedding_model = infer_model('openai:text-embedding-3-small') - -# Use Google's embedding model -embedding_model = infer_model('google-gla:text-embedding-004') - -# Use Cohere for multilingual support -embedding_model = infer_model('cohere:embed-multilingual-v3.0') - -# Pass to VectorGraphMemory -memory = VectorGraphMemory( - qdrant_client=qdrant, - janus_client=janus, - embedding_model=embedding_model -) -``` - -## Complete Example - -```python -from pydantic_ai import Agent, RunContext -from pydantic_ai.models import infer_model -from qdrant_client import QdrantClient -from gremlin_python.driver import client as gremlin_client -from vector_graph_memory import VectorGraphMemory - -# Connect to databases -qdrant = QdrantClient(host='localhost', port=6333) -janus = gremlin_client.Client('ws://localhost:8182/gremlin', 'g') - -# Choose embedding model -embedding_model = infer_model('openai:text-embedding-3-small') - -# Initialize memory system -memory = VectorGraphMemory( - qdrant_client=qdrant, - janus_client=janus, - embedding_model=embedding_model, - collection_name="my_memory" -) - -# Create agent with chosen LLM -agent = Agent('anthropic:claude-3-5-sonnet-20241022') - -# Use different providers for LLM vs embeddings! -# For example: Claude for reasoning + OpenAI for embeddings -``` - -## API Keys - -Set API keys as environment variables: - -```bash -export OPENAI_API_KEY=your-api-key -export ANTHROPIC_API_KEY=your-api-key -export GOOGLE_API_KEY=your-api-key -export COHERE_API_KEY=your-api-key -export VOYAGEAI_API_KEY=your-api-key -``` - -Or in a `.env` file: - -```bash -OPENAI_API_KEY=your-api-key -ANTHROPIC_API_KEY=your-api-key -``` - -## Key Design Principles - -1. **Models are API parameters, not configuration** - Pass model strings directly to `Agent()` and `infer_model()` -2. **Mix and match providers** - Use different providers for LLM vs embeddings (e.g., Claude + OpenAI) -3. **Drop-in replacement** - Change model strings to switch providers without code changes -4. **PydanticAI native** - Leverages PydanticAI's built-in model support - -## Further Reading - -- [PydanticAI Models Documentation](https://ai.pydantic.dev/models/overview/) -- [PydanticAI Embeddings Documentation](https://ai.pydantic.dev/embeddings/) diff --git a/src/vgm/__init__.py b/src/vgm/__init__.py new file mode 100644 index 0000000..28900e1 --- /dev/null +++ b/src/vgm/__init__.py @@ -0,0 +1,54 @@ +"""Vector Graph Memory - Hybrid vector-graph database backend for AI agents.""" + +__version__ = "0.1.0" + +# Main agent +from .MemoryAgent import MemoryAgent + +# Storage layer +from .VectorGraphStore import VectorGraphStore + +# Schemas +from .schemas import ( + NodeMetadata, + EdgeMetadata, + AuditEntry, + SimilarNode, +) + +# Configuration +from .config import ( + MemoryConfig, + MemoryTriggerConfig, + AuditConfig, + VectorGraphConfig, +) + +# Audit backends +from .audit import ( + AuditBackend, + JSONLAuditBackend, + MongoAuditBackend, +) + +__all__ = [ + # Version + "__version__", + # Main API + "MemoryAgent", + "VectorGraphStore", + # Schemas + "NodeMetadata", + "EdgeMetadata", + "AuditEntry", + "SimilarNode", + # Configuration + "MemoryConfig", + "MemoryTriggerConfig", + "AuditConfig", + "VectorGraphConfig", + # Audit backends + "AuditBackend", + "JSONLAuditBackend", + "MongoAuditBackend", +] diff --git a/vector_graph_memory/__init__.py b/vector_graph_memory/__init__.py deleted file mode 100644 index 0d76881..0000000 --- a/vector_graph_memory/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Vector Graph Memory - Hybrid vector-graph database backend for AI agents.""" - -__version__ = "0.1.0" diff --git a/vector_graph_memory/memory.py b/vector_graph_memory/memory.py deleted file mode 100644 index e1270e8..0000000 --- a/vector_graph_memory/memory.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Core memory system integrating vector and graph databases.""" - -import uuid -from typing import List, Dict, Any, Optional -from qdrant_client import QdrantClient -from qdrant_client.models import Distance, VectorParams, PointStruct -from gremlin_python.driver import client as gremlin_client -from pydantic_ai.models import EmbeddingModel - - -class VectorGraphMemory: - """Hybrid vector-graph memory system.""" - - def __init__( - self, - qdrant_client: QdrantClient, - janus_client: gremlin_client.Client, - embedding_model: EmbeddingModel, - collection_name: str = "memory", - ): - """Initialize vector-graph memory system. - - Args: - qdrant_client: Qdrant client instance - janus_client: JanusGraph Gremlin client instance - embedding_model: PydanticAI EmbeddingModel instance - collection_name: Name of the Qdrant collection - """ - self.qdrant = qdrant_client - self.janus = janus_client - self.embedding_model = embedding_model - self.collection_name = collection_name - - self._ensure_collection() - - def _ensure_collection(self) -> None: - """Ensure the Qdrant collection exists.""" - collections = self.qdrant.get_collections().collections - if not any(c.name == self.collection_name for c in collections): - # Get embedding dimension from a test embedding - import asyncio - test_result = asyncio.run(self.embedding_model.embed(["test"])) - embedding_dim = len(test_result.embeddings[0]) - - self.qdrant.create_collection( - collection_name=self.collection_name, - vectors_config=VectorParams( - size=embedding_dim, - distance=Distance.COSINE - ) - ) - - async def add_entity( - self, - content: str, - entity_type: str, - metadata: Optional[Dict[str, Any]] = None - ) -> str: - """Add an entity to both vector and graph databases. - - Args: - content: The content/description of the entity - entity_type: Type of entity (e.g., 'job', 'company', 'person') - metadata: Additional metadata for the entity - - Returns: - Entity ID (UUID) - """ - entity_id = str(uuid.uuid4()) - metadata = metadata or {} - - # Generate embedding using PydanticAI - embedding_result = await self.embedding_model.embed([content]) - embedding = embedding_result.embeddings[0] - - # Store in Qdrant - self.qdrant.upsert( - collection_name=self.collection_name, - points=[ - PointStruct( - id=entity_id, - vector=embedding, - payload={ - "content": content, - "type": entity_type, - **metadata - } - ) - ] - ) - - # Store in JanusGraph - gremlin_query = ( - f"g.addV('{entity_type}')" - f".property('id', '{entity_id}')" - f".property('content', content)" - ) - - # Add metadata properties - for key, value in metadata.items(): - # Escape single quotes in values - escaped_value = str(value).replace("'", "\\'") - gremlin_query += f".property('{key}', '{escaped_value}')" - - self.janus.submit(gremlin_query).all().result() - - return entity_id - - async def search_similar( - self, - query: str, - limit: int = 5, - entity_type: Optional[str] = None - ) -> List[Dict[str, Any]]: - """Search for semantically similar entities. - - Args: - query: Search query text - limit: Maximum number of results - entity_type: Optional filter by entity type - - Returns: - List of matching entities with scores - """ - # Generate query embedding using PydanticAI - embedding_result = await self.embedding_model.embed([query]) - query_embedding = embedding_result.embeddings[0] - - # Search in Qdrant - search_kwargs = { - "collection_name": self.collection_name, - "query_vector": query_embedding, - "limit": limit - } - - if entity_type: - search_kwargs["query_filter"] = { - "must": [{"key": "type", "match": {"value": entity_type}}] - } - - results = self.qdrant.search(**search_kwargs) - - # Format results - return [ - { - "id": result.id, - "score": result.score, - "content": result.payload.get("content"), - "type": result.payload.get("type"), - "metadata": { - k: v for k, v in result.payload.items() - if k not in ["content", "type"] - } - } - for result in results - ] - - async def add_relationship( - self, - from_id: str, - to_id: str, - relationship_type: str, - properties: Optional[Dict[str, Any]] = None - ) -> None: - """Add a relationship between two entities in the graph. - - Args: - from_id: Source entity ID - to_id: Target entity ID - relationship_type: Type of relationship (e.g., 'works_at', 'applied_to') - properties: Optional properties for the relationship - """ - properties = properties or {} - - # Build Gremlin query to add edge - gremlin_query = ( - f"g.V().has('id', '{from_id}')" - f".addE('{relationship_type}')" - f".to(g.V().has('id', '{to_id}'))" - ) - - # Add properties to edge - for key, value in properties.items(): - escaped_value = str(value).replace("'", "\\'") - gremlin_query += f".property('{key}', '{escaped_value}')" - - self.janus.submit(gremlin_query).all().result() - - async def query_graph_traversal( - self, - start_id: str, - gremlin_steps: str - ) -> List[Dict[str, Any]]: - """Execute a graph traversal starting from a specific entity. - - Args: - start_id: Starting entity ID - gremlin_steps: Gremlin traversal steps to execute - - Returns: - List of results from the traversal - """ - gremlin_query = f"g.V().has('id', '{start_id}').{gremlin_steps}" - results = self.janus.submit(gremlin_query).all().result() - return results From 1bd7a57120f5cc25249489d0a1a3f2f16495becd Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:09 -0700 Subject: [PATCH 13/21] feat: implement core vector-graph memory system MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add VectorGraphStore and MemoryAgent for hybrid memory: - VectorGraphStore: Qdrant (vectors) + JanusGraph (graph) integration - MemoryAgent: PydanticAI agent with memory proposal workflow - Schemas: Pydantic models for nodes, edges, proposals, audit - Memory triggers: phrase, interval, and AI-determined modes - Similarity-based duplicate detection 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- src/vgm/MemoryAgent.py | 496 ++++++++++++++++++++++++++++++++ src/vgm/VectorGraphStore.py | 475 ++++++++++++++++++++++++++++++ src/vgm/schemas/AuditEntry.py | 18 ++ src/vgm/schemas/EdgeMetadata.py | 21 ++ src/vgm/schemas/NodeMetadata.py | 20 ++ src/vgm/schemas/SimilarNode.py | 14 + src/vgm/schemas/__init__.py | 13 + 7 files changed, 1057 insertions(+) create mode 100644 src/vgm/MemoryAgent.py create mode 100644 src/vgm/VectorGraphStore.py create mode 100644 src/vgm/schemas/AuditEntry.py create mode 100644 src/vgm/schemas/EdgeMetadata.py create mode 100644 src/vgm/schemas/NodeMetadata.py create mode 100644 src/vgm/schemas/SimilarNode.py create mode 100644 src/vgm/schemas/__init__.py diff --git a/src/vgm/MemoryAgent.py b/src/vgm/MemoryAgent.py new file mode 100644 index 0000000..6b08551 --- /dev/null +++ b/src/vgm/MemoryAgent.py @@ -0,0 +1,496 @@ +"""Memory-enhanced AI agent with vector-graph storage.""" + +import asyncio +from typing import Optional, List, Dict, Any, Literal +from uuid import uuid4 + +from pydantic_ai import Agent, RunContext +from pydantic_ai.models import Model +from pydantic import BaseModel +from qdrant_client import QdrantClient +from gremlin_python.driver import client as gremlin_client +from pydantic_ai import EmbeddingModel + +from .VectorGraphStore import VectorGraphStore, _run_async +from .schemas import NodeMetadata, EdgeMetadata, AuditEntry +from .config import MemoryConfig, MemoryTriggerConfig, AuditConfig, VectorGraphConfig +from .audit import AuditBackend, JSONLAuditBackend, MongoAuditBackend + + +class AgentDependencies(BaseModel): + """Dependencies passed to agent tools.""" + + memory_agent: Any # Reference to MemoryAgent instance + + class Config: + arbitrary_types_allowed = True + + +class MemoryAgent: + """AI agent with persistent vector-graph memory. + + This agent wraps PydanticAI with memory capabilities, allowing it to: + - Search semantic memory via vector similarity + - Store entities and relationships in a graph + - Track memory operations via audit logs + - Propose memory additions with user confirmation + """ + + def __init__( + self, + qdrant_client: QdrantClient, + janus_client: gremlin_client.Client, + embedding_model: EmbeddingModel, + llm_model: str | Model, + system_prompt: str, + memory_config: MemoryConfig, + audit_config: Optional[AuditConfig] = None, + vector_config: Optional[VectorGraphConfig] = None, + trigger_config: Optional[MemoryTriggerConfig] = None, + ): + """Initialize memory-enhanced agent. + + Args: + qdrant_client: Qdrant client instance + janus_client: JanusGraph Gremlin client instance + embedding_model: PydanticAI EmbeddingModel for embeddings + llm_model: LLM model identifier (e.g., 'openai:gpt-4') + system_prompt: Base system prompt for the agent + memory_config: Memory behavior configuration + audit_config: Audit logging configuration (optional) + vector_config: Vector/graph database configuration (optional) + trigger_config: Memory trigger configuration (optional) + """ + self.memory_config = memory_config + self.trigger_config = trigger_config or MemoryTriggerConfig(mode="phrase", trigger_phrase="check for memory items") + self.audit_config = audit_config or AuditConfig() + + # Initialize storage layer + self.store = VectorGraphStore( + qdrant_client=qdrant_client, + janus_client=janus_client, + embedding_model=embedding_model, + config=vector_config, + ) + + # Initialize audit backend + self.audit = self._create_audit_backend(self.audit_config) + + # Session tracking + self._current_session_id: Optional[str] = None + self._message_count: int = 0 + + # Pending memory proposals (for user confirmation) + self.pending_proposals: Dict[str, Dict[str, Any]] = {} + + # Build system prompt with memory instructions + full_system_prompt = self._build_system_prompt(system_prompt) + + # Create PydanticAI agent with tools + self.agent = Agent( + llm_model, + deps_type=AgentDependencies, + system_prompt=full_system_prompt, + ) + + # Register memory tools + self._register_tools() + + @staticmethod + def _create_audit_backend(config: AuditConfig) -> AuditBackend: + """Create appropriate audit backend based on config.""" + if config.backend == "jsonl": + return JSONLAuditBackend(config) + elif config.backend == "mongodb": + return MongoAuditBackend(config) + else: + raise ValueError(f"Unknown audit backend: {config.backend}") + + def _build_system_prompt(self, base_prompt: str) -> str: + """Build full system prompt with memory instructions.""" + memory_instructions = f""" + +## Memory System + +You have access to a vector-graph memory system for persistent storage. + +**Use Case:** {self.memory_config.use_case_description} + +**Memory Threshold:** {self.memory_config.memory_threshold_description} + +**Available Memory Tools:** +- `search_memory(query, limit)` - Find semantically similar content in memory +- `propose_memory_addition(content, entity_type, relationships)` - Suggest storing new information +- `get_memory_context(node_id, traversal_steps)` - Explore relationships via graph traversal + +When asked about your capabilities or "what tools" you have, explain these memory management tools. + +**Important Rules:** +1. Before proposing new memory, ALWAYS search for similar existing content first +2. When proposing memory, ask the user: "Should I add [description] to memory?" +3. NEVER add to memory without explicit user confirmation +4. After successful memory operation, provide a brief status: "Successfully added [item] to memory" +5. If similar content exists, ask if the user wants to update existing or create new + +**Project ID:** {self.memory_config.project_id} +""" + return base_prompt + memory_instructions + + def _register_tools(self) -> None: + """Register memory tools with the agent.""" + + @self.agent.tool + def search_memory( + ctx: RunContext[AgentDependencies], + query: str, + limit: int = 5, + ) -> str: + """Search the memory for semantically similar content. + + Args: + ctx: Agent runtime context + query: Natural language search query + limit: Maximum number of results (default: 5) + + Returns: + Formatted search results + """ + results = ctx.deps.memory_agent.store.search_similar_nodes( + content=query, + limit=limit, + project_id=ctx.deps.memory_agent.memory_config.project_id, + ) + + if not results: + return f"No similar content found in memory for: '{query}'" + + formatted = [] + for i, result in enumerate(results, 1): + formatted.append( + f"{i}. [{result.node_type}] {result.content[:200]}... " + f"(similarity: {result.similarity_score:.2f})" + ) + + return "\n".join(formatted) + + @self.agent.tool + def propose_memory_addition( + ctx: RunContext[AgentDependencies], + content: str, + entity_type: str, + relationships: Optional[List[Dict[str, str]]] = None, + ) -> str: + """Propose adding content to memory (requires user confirmation). + + Args: + ctx: Agent runtime context + content: The content to store + entity_type: Type of entity (e.g., 'job', 'company', 'person') + relationships: Optional list of relationships to create + Each relationship should have: {"to_node_id": "...", "type": "...", "description": "..."} + + Returns: + Proposal message for user confirmation + """ + agent_instance = ctx.deps.memory_agent + + # Search for duplicates + similar = agent_instance.store.search_similar_nodes( + content=content, + threshold=agent_instance.memory_config.similarity_threshold, + limit=3, + project_id=agent_instance.memory_config.project_id, + ) + + # Generate proposal ID + proposal_id = str(uuid4()) + + # Store proposal for later execution + agent_instance.pending_proposals[proposal_id] = { + "content": content, + "entity_type": entity_type, + "relationships": relationships or [], + "similar_nodes": similar, + } + + # Build proposal message + if similar: + proposal = f"Found {len(similar)} similar item(s) in memory:\n" + for i, sim in enumerate(similar, 1): + proposal += f" {i}. [{sim.node_type}] {sim.content[:100]}... (similarity: {sim.similarity_score:.2f})\n" + proposal += f"\nShould I add '{content[:100]}...' as a new {entity_type} to memory, or update one of the existing items?" + else: + proposal = f"Should I add '{content[:100]}...' as a {entity_type} to memory?" + + return proposal + + @self.agent.tool + def get_memory_context( + ctx: RunContext[AgentDependencies], + node_id: str, + traversal_steps: str = "both().limit(10)", + ) -> str: + """Get context around a memory node via graph traversal. + + Args: + ctx: Agent runtime context + node_id: Starting node ID + traversal_steps: Gremlin traversal steps (default: "both().limit(10)") + + Returns: + Formatted traversal results + """ + try: + results = ctx.deps.memory_agent.store.traverse_from_node( + node_id=node_id, + gremlin_steps=traversal_steps, + ) + + if not results: + return f"No connected nodes found for {node_id}" + + formatted = [] + for i, node in enumerate(results, 1): + if isinstance(node, dict) and "content" in node: + formatted.append(f"{i}. [{node.get('node_type', 'unknown')}] {node['content'][:200]}...") + else: + formatted.append(f"{i}. {node}") + + return "\n".join(formatted) + except Exception as e: + return f"Error traversing from node {node_id}: {str(e)}" + + def configure_memory( + self, + use_case_description: Optional[str] = None, + memory_threshold_description: Optional[str] = None, + similarity_threshold: Optional[float] = None, + ) -> None: + """Update memory configuration dynamically. + + Args: + use_case_description: New use case description + memory_threshold_description: New memory threshold description + similarity_threshold: New similarity threshold for duplicate detection + """ + if use_case_description: + self.memory_config.use_case_description = use_case_description + if memory_threshold_description: + self.memory_config.memory_threshold_description = memory_threshold_description + if similarity_threshold is not None: + self.memory_config.similarity_threshold = similarity_threshold + + # Rebuild system prompt + # Note: This updates the config but doesn't re-create the agent + # The agent will use the new config on next run + + def set_memory_trigger( + self, + mode: Literal["phrase", "interval", "ai_determined"], + trigger_phrase: Optional[str] = None, + message_interval: Optional[int] = None, + ) -> None: + """Update memory trigger configuration dynamically. + + Args: + mode: Trigger mode ('phrase', 'interval', 'ai_determined') + trigger_phrase: Phrase to trigger memory check (for 'phrase' mode) + message_interval: Number of messages between checks (for 'interval' mode) + """ + self.trigger_config = MemoryTriggerConfig( + mode=mode, + trigger_phrase=trigger_phrase, + message_interval=message_interval, + ) + + def confirm_memory_addition( + self, + proposal_id: str, + action: str = "add_new", + update_node_id: Optional[str] = None, + ) -> str: + """Confirm and execute a pending memory proposal. + + Args: + proposal_id: ID of the pending proposal + action: Action to take ('add_new', 'update_existing', 'cancel') + update_node_id: Node ID to update (required if action='update_existing') + + Returns: + Status message + """ + if proposal_id not in self.pending_proposals: + return f"No pending proposal found with ID: {proposal_id}" + + proposal = self.pending_proposals[proposal_id] + + if action == "cancel": + del self.pending_proposals[proposal_id] + return "Memory addition cancelled." + + try: + if action == "add_new": + # Create new node + node_metadata = NodeMetadata( + content=proposal["content"], + node_type=proposal["entity_type"], + source=self._current_session_id or "unknown", + project_id=self.memory_config.project_id, + embedding_model=str(self.store.embedding_model), + ) + + node_id = self.store.add_node(node_metadata) + + # Add relationships if specified + edge_ids = [] + for rel in proposal["relationships"]: + edge_metadata = EdgeMetadata( + from_node_id=node_id, + to_node_id=rel["to_node_id"], + relationship_type=rel["type"], + description=rel.get("description", ""), + source=self._current_session_id or "unknown", + project_id=self.memory_config.project_id, + ) + edge_id = self.store.add_edge(edge_metadata) + edge_ids.append(edge_id) + + # Log to audit + audit_entry = AuditEntry( + session_id=self._current_session_id or "unknown", + project_id=self.memory_config.project_id, + operation_type="add_node", + summary=f"Added {proposal['entity_type']}: {proposal['content'][:100]}", + commands=[f"add_node(node_id={node_id})", *[f"add_edge(edge_id={eid})" for eid in edge_ids]], + metadata={"node_type": proposal["entity_type"], "relationship_count": len(edge_ids)}, + affected_entities=[node_id, *edge_ids], + ) + self.audit.log_operation(audit_entry) + + del self.pending_proposals[proposal_id] + return f"Successfully added {proposal['entity_type']} to memory with ID: {node_id}" + + elif action == "update_existing": + if not update_node_id: + return "Error: update_node_id required for update_existing action" + + # Update existing node + self.store.update_node( + node_id=update_node_id, + content=proposal["content"], + ) + + # Log to audit + audit_entry = AuditEntry( + session_id=self._current_session_id or "unknown", + project_id=self.memory_config.project_id, + operation_type="update_node", + summary=f"Updated {proposal['entity_type']}: {proposal['content'][:100]}", + commands=[f"update_node(node_id={update_node_id})"], + metadata={"node_type": proposal["entity_type"]}, + affected_entities=[update_node_id], + ) + self.audit.log_operation(audit_entry) + + del self.pending_proposals[proposal_id] + return f"Successfully updated existing node: {update_node_id}" + + else: + return f"Unknown action: {action}" + + except Exception as e: + import traceback + error_details = traceback.format_exc() + return f"Error executing memory operation: {str(e)}\n\nFull traceback:\n{error_details}" + + def run( + self, + prompt: str, + session_id: Optional[str] = None, + ) -> Any: + """Run the agent with a user prompt. + + Args: + prompt: User input + session_id: Optional session ID (auto-generated if not provided) + + Returns: + Agent response + """ + # Set/update session ID + if session_id: + if session_id != self._current_session_id: + # New session, reset message counter + self._current_session_id = session_id + self._message_count = 0 + else: + if not self._current_session_id: + self._current_session_id = str(uuid4()) + + # Increment message count + self._message_count += 1 + + # Check if we should trigger memory review + should_check_memory = self._should_check_memory(prompt) + + # Optionally inject memory check instruction + if should_check_memory: + enhanced_prompt = f"{prompt}\n\n[Assistant instruction: Review this conversation for items worth adding to memory based on the configured memory threshold. Use the memory tools to propose additions if appropriate.]" + else: + enhanced_prompt = prompt + + # Create dependencies + deps = AgentDependencies(memory_agent=self) + + # Run agent (synchronous for v1) + result = _run_async(self.agent.run(enhanced_prompt, deps=deps)) + + return result + + def _should_check_memory(self, prompt: str) -> bool: + """Determine if we should check for memory items. + + Args: + prompt: User prompt + + Returns: + True if memory check should be triggered + """ + if self.trigger_config.mode == "phrase": + if self.trigger_config.trigger_phrase: + return self.trigger_config.trigger_phrase.lower() in prompt.lower() + return False + + elif self.trigger_config.mode == "interval": + if self.trigger_config.message_interval: + return self._message_count % self.trigger_config.message_interval == 0 + return False + + elif self.trigger_config.mode == "ai_determined": + # Always let AI decide + return True + + return False + + def get_audit_history( + self, + session_id: Optional[str] = None, + entity_id: Optional[str] = None, + limit: int = 50, + ) -> List[AuditEntry]: + """Get audit log history. + + Args: + session_id: Filter by session ID + entity_id: Filter by entity ID + limit: Maximum number of entries + + Returns: + List of audit entries + """ + if session_id: + return self.audit.get_by_session(session_id) + elif entity_id: + return self.audit.get_entity_history(entity_id) + else: + return self.audit.get_recent(limit) diff --git a/src/vgm/VectorGraphStore.py b/src/vgm/VectorGraphStore.py new file mode 100644 index 0000000..61ec4b5 --- /dev/null +++ b/src/vgm/VectorGraphStore.py @@ -0,0 +1,475 @@ +"""Vector and graph database storage layer.""" + +import asyncio +from typing import List, Optional, Dict, Any, Coroutine, TypeVar +from datetime import datetime, timezone + +from qdrant_client import QdrantClient +from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue +from gremlin_python.driver import client as gremlin_client +from pydantic_ai import EmbeddingModel + +from .schemas import NodeMetadata, EdgeMetadata, SimilarNode +from .config import VectorGraphConfig + +T = TypeVar('T') + + +def _run_async(coro: Coroutine[Any, Any, T]) -> T: + """Run async code, handling both notebook and regular Python environments. + + This helper manages event loop differences between: + - Regular Python (no running loop) + - Jupyter notebooks (with running loop) + - FastAPI/uvicorn (with uvloop that can't be patched) + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # No running loop - we're in regular Python + return asyncio.run(coro) + else: + # There's a running loop + # Check if it's patchable (uvloop isn't, but asyncio default loop is) + import nest_asyncio + import threading + loop_type = type(loop).__name__ + + if loop_type == 'Loop' and 'uvloop' in loop.__class__.__module__: + # uvloop can't be patched - run in new thread with new loop + result = None + exception = None + + def run_in_thread(): + nonlocal result, exception + try: + result = asyncio.run(coro) + except Exception as e: + exception = e + + thread = threading.Thread(target=run_in_thread) + thread.start() + thread.join() + + if exception: + raise exception + return result + else: + # Regular asyncio loop - can be patched for notebooks + nest_asyncio.apply() + return asyncio.run(coro) + + +def _run_in_thread(func, *args, **kwargs): + """Run a synchronous function in a separate thread. + + This is needed for Gremlin client operations which internally use + run_until_complete() and fail when called from an existing event loop + (like FastAPI/uvicorn). + """ + import threading + result = None + exception = None + + def run(): + nonlocal result, exception + try: + result = func(*args, **kwargs) + except Exception as e: + exception = e + + thread = threading.Thread(target=run) + thread.start() + thread.join() + + if exception: + raise exception + return result + + +class VectorGraphStore: + """Hybrid vector-graph memory storage system. + + Storage strategy: + - Qdrant: Vectors + full content + metadata (source of truth for data) + - JanusGraph: Node references + relationships only (source of truth for structure) + """ + + # Content size limit (1 MB soft limit, configurable) + DEFAULT_MAX_CONTENT_SIZE = 1_000_000 # ~1 million characters + + def __init__( + self, + qdrant_client: QdrantClient, + janus_client: gremlin_client.Client, + embedding_model: EmbeddingModel, + config: Optional[VectorGraphConfig] = None, + max_content_size: Optional[int] = None, + ): + """Initialize vector-graph store. + + Args: + qdrant_client: Qdrant client instance + janus_client: JanusGraph Gremlin client instance + embedding_model: PydanticAI EmbeddingModel instance + config: Storage configuration (optional) + max_content_size: Maximum content size in characters (default: 1MB) + """ + self.qdrant = qdrant_client + self.janus = janus_client + self.embedding_model = embedding_model + self.config = config or VectorGraphConfig() + self.max_content_size = max_content_size or self.DEFAULT_MAX_CONTENT_SIZE + + self._ensure_janusgraph_schema() + self._ensure_collection() + + def _ensure_janusgraph_schema(self) -> None: + """Check if JanusGraph schema is initialized. + + Note: Schema initialization is skipped in this method to avoid asyncio + event loop conflicts in Jupyter notebooks. Run scripts/init_janusgraph_schema.py + once before using the library. + """ + # Skip automatic schema initialization - must be done manually + # This avoids asyncio event loop conflicts in Jupyter notebooks + pass + + @staticmethod + def _escape_gremlin_value(val: Any) -> str: + """Escape values for safe use in Gremlin queries.""" + return str(val).replace("'", "\\'").replace('"', '\\"') + + def _ensure_collection(self) -> None: + """Ensure the Qdrant collection exists.""" + collections = self.qdrant.get_collections().collections + if not any(c.name == self.config.qdrant_collection for c in collections): + # Get embedding dimension from a test embedding + test_result = _run_async(self.embedding_model.embed(["test"], input_type="document")) + embedding_dim = len(test_result.embeddings[0]) + + self.qdrant.create_collection( + collection_name=self.config.qdrant_collection, + vectors_config=VectorParams( + size=embedding_dim, + distance=Distance.COSINE + ) + ) + + def _generate_embedding(self, text: str) -> List[float]: + """Generate embedding for text using the embedding model. + + Args: + text: Text to embed + + Returns: + Embedding vector + """ + result = _run_async(self.embedding_model.embed([text], input_type="document")) + return list(result.embeddings[0]) + + def search_similar_nodes( + self, + content: str, + threshold: Optional[float] = None, + limit: int = 5, + project_id: Optional[str] = None, + ) -> List[SimilarNode]: + """Search for similar existing nodes using Top-K retrieval. + + Args: + content: Text to search for + threshold: Optional minimum similarity score (0.0-1.0). If None, returns top K results. + limit: Maximum number of results (default: 5) + project_id: Filter by project (optional) + + Returns: + List of top K similar nodes, optionally filtered by threshold + """ + import logging + logger = logging.getLogger(__name__) + + embedding = self._generate_embedding(content) + + # Build search parameters + query_filter = None + if project_id: + query_filter = Filter( + must=[FieldCondition(key="project_id", match=MatchValue(value=project_id))] + ) + + logger.info(f"[SEARCH] Query: '{content[:100]}...'") + logger.info(f"[SEARCH] Project ID: {project_id}") + logger.info(f"[SEARCH] Threshold: {threshold}, Limit: {limit}") + + response = self.qdrant.query_points( + collection_name=self.config.qdrant_collection, + query=embedding, # type: ignore[arg-type] # list[float] is valid but not in type stubs + limit=limit, + score_threshold=threshold, # None means no threshold - Top-K retrieval + query_filter=query_filter, + ) + + logger.info(f"[SEARCH] Found {len(response.points)} results") + for i, result in enumerate(response.points, 1): + logger.info(f"[SEARCH] {i}. [{result.payload.get('node_type')}] score={result.score:.3f} content='{str(result.payload.get('content', ''))[:80]}...'") + + return [ + SimilarNode( + node_id=str(result.id), + content=result.payload["content"], + node_type=result.payload["node_type"], + similarity_score=result.score, + metadata={k: v for k, v in result.payload.items() if k not in ["content", "node_type"]}, + ) + for result in response.points + ] + + def add_node( + self, + metadata: NodeMetadata, + ) -> str: + """Add node to both vector and graph databases. + + Args: + metadata: Complete node metadata including content + + Returns: + Node ID (UUID) + + Raises: + ValueError: If content exceeds max_content_size + """ + import logging + logger = logging.getLogger(__name__) + + # Check content size + if len(metadata.content) > self.max_content_size: + raise ValueError( + f"Content size ({len(metadata.content)} chars) exceeds maximum " + f"({self.max_content_size} chars). Consider chunking large content." + ) + + # Generate embedding + embedding = self._generate_embedding(metadata.content) + + # Store FULL DATA in Qdrant (source of truth for content) + payload = { + "content": metadata.content, + "node_type": metadata.node_type, + "created_at": metadata.created_at.isoformat(), + "updated_at": metadata.updated_at.isoformat(), + "source": metadata.source, + "project_id": metadata.project_id, + "embedding_model": metadata.embedding_model, + **metadata.custom_metadata, + } + + logger.info(f"[STORE] Adding node ID: {metadata.node_id}") + logger.info(f"[STORE] Type: {metadata.node_type}") + logger.info(f"[STORE] Project ID: {metadata.project_id}") + logger.info(f"[STORE] Content: '{metadata.content[:100]}...'") + + self.qdrant.upsert( + collection_name=self.config.qdrant_collection, + points=[ + PointStruct( + id=metadata.node_id, + vector=embedding, + payload=payload, + ) + ], + ) + + logger.info(f"[STORE] Successfully added to Qdrant") + + # Store MINIMAL REFERENCE in JanusGraph (structure only) + # Use generic "node" label, store actual type as a property + # Use iterate() to avoid returning the vertex (avoids serialization issues) + gremlin_query = ( + f"g.addV('node')" + f".property('node_id', '{self._escape_gremlin_value(metadata.node_id)}')" + f".property('node_type', '{self._escape_gremlin_value(metadata.node_type)}')" + f".property('project_id', '{self._escape_gremlin_value(metadata.project_id)}')" + f".property('created_at', '{metadata.created_at.isoformat()}')" + f".property('updated_at', '{metadata.updated_at.isoformat()}')" + f".iterate()" + ) + _run_in_thread(lambda: self.janus.submit(gremlin_query).all().result()) + + return metadata.node_id + + def update_node( + self, + node_id: str, + content: Optional[str] = None, + custom_metadata: Optional[Dict[str, Any]] = None, + ) -> None: + """Update existing node. + + Args: + node_id: Node ID to update + content: New content (will regenerate embedding if provided) + custom_metadata: Metadata fields to update + + Raises: + ValueError: If node not found or content exceeds max_content_size + """ + # Get current node from Qdrant + current = self.qdrant.retrieve( + collection_name=self.config.qdrant_collection, + ids=[node_id], + ) + + if not current: + raise ValueError(f"Node {node_id} not found") + + current_payload = current[0].payload + updated_at = datetime.now(timezone.utc).isoformat() + + # Update payload + if content: + # Check content size + if len(content) > self.max_content_size: + raise ValueError( + f"Content size ({len(content)} chars) exceeds maximum " + f"({self.max_content_size} chars)" + ) + + current_payload["content"] = content + current_payload["updated_at"] = updated_at + # Regenerate embedding + embedding = self._generate_embedding(content) + else: + # Keep existing embedding + embedding = self.qdrant.retrieve( + collection_name=self.config.qdrant_collection, + ids=[node_id], + with_vectors=True, + )[0].vector + + if custom_metadata: + current_payload.update(custom_metadata) + current_payload["updated_at"] = updated_at + + # Update in Qdrant (content + vector) + self.qdrant.upsert( + collection_name=self.config.qdrant_collection, + points=[ + PointStruct( + id=node_id, + vector=embedding, + payload=current_payload, + ) + ], + ) + + # Update timestamp in JanusGraph (structure only) + update_query = ( + f"g.V().has('node_id', '{self._escape_gremlin_value(node_id)}')" + f".property('updated_at', '{updated_at}')" + ) + + _run_in_thread(lambda: self.janus.submit(update_query).all().result()) + + def add_edge( + self, + metadata: EdgeMetadata, + ) -> str: + """Add edge to graph database. + + Args: + metadata: Complete edge metadata + + Returns: + Edge ID (UUID) + """ + # Build Gremlin query to add edge + # Use generic "relationship" label, store actual type as property + # Use iterate() to avoid returning the edge (avoids serialization issues) + gremlin_query = ( + f"g.V().has('node_id', '{self._escape_gremlin_value(metadata.from_node_id)}')" + f".addE('relationship')" + f".to(g.V().has('node_id', '{self._escape_gremlin_value(metadata.to_node_id)}'))" + f".property('edge_id', '{self._escape_gremlin_value(metadata.edge_id)}')" + f".property('relationship_type', '{self._escape_gremlin_value(metadata.relationship_type)}')" + f".property('description', '{self._escape_gremlin_value(metadata.description)}')" + f".property('created_at', '{metadata.created_at.isoformat()}')" + f".property('source', '{self._escape_gremlin_value(metadata.source)}')" + f".property('project_id', '{self._escape_gremlin_value(metadata.project_id)}')" + ) + if metadata.confidence is not None: + gremlin_query += f".property('confidence', {metadata.confidence})" + for key, value in metadata.custom_metadata.items(): + gremlin_query += f".property('{self._escape_gremlin_value(key)}', '{self._escape_gremlin_value(value)}')" + gremlin_query += ".iterate()" + _run_in_thread(lambda: self.janus.submit(gremlin_query).all().result()) + + return metadata.edge_id + + def get_node(self, node_id: str) -> Optional[Dict[str, Any]]: + """Retrieve node by ID from Qdrant. + + Args: + node_id: Node ID + + Returns: + Node data (including content) or None if not found + """ + results = self.qdrant.retrieve( + collection_name=self.config.qdrant_collection, + ids=[node_id], + ) + + if not results: + return None + + return results[0].payload + + def get_nodes_batch(self, node_ids: List[str]) -> Dict[str, Dict[str, Any]]: + """Retrieve multiple nodes by ID from Qdrant. + + Args: + node_ids: List of node IDs + + Returns: + Dictionary mapping node_id to node data + """ + results = self.qdrant.retrieve( + collection_name=self.config.qdrant_collection, + ids=node_ids, + ) + + return {str(result.id): result.payload for result in results} + + def traverse_from_node( + self, + node_id: str, + gremlin_steps: str, + ) -> List[Dict[str, Any]]: + """Execute a graph traversal starting from a specific node. + + Args: + node_id: Starting node ID + gremlin_steps: Gremlin traversal steps to execute + + Returns: + List of traversal results with full node content + + Note: + If traversal returns node IDs, this will batch-fetch full content from Qdrant + """ + gremlin_query = ( + f"g.V().has('node_id', '{self._escape_gremlin_value(node_id)}')" + f".{gremlin_steps}" + ) + results = _run_in_thread(lambda: self.janus.submit(gremlin_query).all().result()) + + # If results contain node_ids, fetch full content from Qdrant + # This is a simple heuristic - results might be node_ids or other data + if results and all(isinstance(r, str) for r in results): + # Looks like node IDs, batch fetch from Qdrant + node_data = self.get_nodes_batch(results) + return [node_data.get(node_id, {"node_id": node_id}) for node_id in results] + + return results diff --git a/src/vgm/schemas/AuditEntry.py b/src/vgm/schemas/AuditEntry.py new file mode 100644 index 0000000..4594112 --- /dev/null +++ b/src/vgm/schemas/AuditEntry.py @@ -0,0 +1,18 @@ +"""Audit entry schema for memory operations.""" + +from datetime import datetime +from typing import Any, Dict, List, Literal +from pydantic import BaseModel, Field + + +class AuditEntry(BaseModel): + """Log entry for memory operations.""" + + timestamp: datetime = Field(default_factory=datetime.utcnow) + session_id: str + project_id: str + operation_type: Literal["add_node", "update_node", "add_edge", "merge_nodes"] + summary: str # Human-readable description of what happened + commands: List[str] # Database commands executed (content sanitized) + metadata: Dict[str, Any] # Metadata that was set (no content) + affected_entities: List[str] # Node/edge IDs involved diff --git a/src/vgm/schemas/EdgeMetadata.py b/src/vgm/schemas/EdgeMetadata.py new file mode 100644 index 0000000..c61ee4f --- /dev/null +++ b/src/vgm/schemas/EdgeMetadata.py @@ -0,0 +1,21 @@ +"""Edge metadata schema for graph relationships.""" + +from datetime import datetime +from typing import Any, Dict, Optional +from pydantic import BaseModel, Field +from uuid import uuid4 + + +class EdgeMetadata(BaseModel): + """Metadata for graph edges/relationships.""" + + edge_id: str = Field(default_factory=lambda: str(uuid4())) + from_node_id: str + to_node_id: str + relationship_type: str # e.g., "works_at", "located_in", "applied_to" + description: str = "" # Freeform context about the relationship + created_at: datetime = Field(default_factory=datetime.utcnow) + source: str # session_id that created this edge + project_id: str + confidence: Optional[float] = Field(default=None, ge=0.0, le=1.0) # Agent confidence in relationship + custom_metadata: Dict[str, Any] = Field(default_factory=dict) diff --git a/src/vgm/schemas/NodeMetadata.py b/src/vgm/schemas/NodeMetadata.py new file mode 100644 index 0000000..334b465 --- /dev/null +++ b/src/vgm/schemas/NodeMetadata.py @@ -0,0 +1,20 @@ +"""Node metadata schema for graph nodes.""" + +from datetime import datetime +from typing import Any, Dict +from pydantic import BaseModel, Field +from uuid import uuid4 + + +class NodeMetadata(BaseModel): + """Metadata for graph nodes stored in both vector and graph databases.""" + + node_id: str = Field(default_factory=lambda: str(uuid4())) + node_type: str # e.g., "job", "company", "person", "interaction" + content: str # The actual text content to be embedded + created_at: datetime = Field(default_factory=datetime.utcnow) + updated_at: datetime = Field(default_factory=datetime.utcnow) + source: str # session_id that created this node + project_id: str # which project/memory space this belongs to + embedding_model: str # e.g., "text-embedding-3-small" + custom_metadata: Dict[str, Any] = Field(default_factory=dict) # User-defined fields diff --git a/src/vgm/schemas/SimilarNode.py b/src/vgm/schemas/SimilarNode.py new file mode 100644 index 0000000..f17e955 --- /dev/null +++ b/src/vgm/schemas/SimilarNode.py @@ -0,0 +1,14 @@ +"""Similar node schema for duplicate detection.""" + +from typing import Any, Dict +from pydantic import BaseModel, Field + + +class SimilarNode(BaseModel): + """Result from similarity search for duplicate detection.""" + + node_id: str + content: str + node_type: str + similarity_score: float = Field(ge=0.0, le=1.0) + metadata: Dict[str, Any] diff --git a/src/vgm/schemas/__init__.py b/src/vgm/schemas/__init__.py new file mode 100644 index 0000000..e673a77 --- /dev/null +++ b/src/vgm/schemas/__init__.py @@ -0,0 +1,13 @@ +"""Schema exports for Vector Graph Memory.""" + +from .NodeMetadata import NodeMetadata +from .EdgeMetadata import EdgeMetadata +from .AuditEntry import AuditEntry +from .SimilarNode import SimilarNode + +__all__ = [ + "NodeMetadata", + "EdgeMetadata", + "AuditEntry", + "SimilarNode", +] From f1bf5ec995fd07b090b338b3708e874e0e2991e3 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:15 -0700 Subject: [PATCH 14/21] feat: add configuration system for memory behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add configurable settings for: - Memory use case and threshold descriptions - Similarity thresholds for duplicate detection - Project-based memory isolation - Memory trigger modes (phrase/interval/AI-determined) - Audit backend selection (JSONL/MongoDB) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- src/vgm/config/AuditConfig.py | 21 +++++++++++++++++++++ src/vgm/config/MemoryConfig.py | 12 ++++++++++++ src/vgm/config/MemoryTriggerConfig.py | 16 ++++++++++++++++ src/vgm/config/VectorGraphConfig.py | 10 ++++++++++ src/vgm/config/__init__.py | 13 +++++++++++++ 5 files changed, 72 insertions(+) create mode 100644 src/vgm/config/AuditConfig.py create mode 100644 src/vgm/config/MemoryConfig.py create mode 100644 src/vgm/config/MemoryTriggerConfig.py create mode 100644 src/vgm/config/VectorGraphConfig.py create mode 100644 src/vgm/config/__init__.py diff --git a/src/vgm/config/AuditConfig.py b/src/vgm/config/AuditConfig.py new file mode 100644 index 0000000..c8f1fe9 --- /dev/null +++ b/src/vgm/config/AuditConfig.py @@ -0,0 +1,21 @@ +"""Audit configuration schema.""" + +from typing import Literal, Optional +from pydantic import BaseModel, Field + + +class AuditConfig(BaseModel): + """Configuration for audit log backend.""" + + backend: Literal["jsonl", "mongodb"] = "jsonl" + + # JSONL-specific options + log_dir: Optional[str] = Field(default=None) # Default: ~/.vgm/logs + rotation_size_mb: int = Field(default=10, gt=0) + rotation_period: Literal["daily", "weekly", "monthly"] = "monthly" + + # MongoDB-specific options + connection_string: Optional[str] = Field(default=None) # e.g., "mongodb://localhost:27017" + database: str = "vgm_audit" + collection: str = "logs" + ttl_days: Optional[int] = Field(default=None, gt=0) # Auto-expire old logs diff --git a/src/vgm/config/MemoryConfig.py b/src/vgm/config/MemoryConfig.py new file mode 100644 index 0000000..131e602 --- /dev/null +++ b/src/vgm/config/MemoryConfig.py @@ -0,0 +1,12 @@ +"""Memory configuration schema.""" + +from pydantic import BaseModel, Field + + +class MemoryConfig(BaseModel): + """Configuration for memory behavior and thresholds.""" + + use_case_description: str # Natural language: "Track job search activities..." + memory_threshold_description: str # Natural language: "Store specific job postings..." + project_id: str # Unique identifier for this memory project + similarity_threshold: float = Field(default=0.85, ge=0.0, le=1.0) # Duplicate detection threshold diff --git a/src/vgm/config/MemoryTriggerConfig.py b/src/vgm/config/MemoryTriggerConfig.py new file mode 100644 index 0000000..bf1f2d7 --- /dev/null +++ b/src/vgm/config/MemoryTriggerConfig.py @@ -0,0 +1,16 @@ +"""Memory trigger configuration schema.""" + +from typing import Literal, Optional +from pydantic import BaseModel, Field + + +class MemoryTriggerConfig(BaseModel): + """Configuration for when to check conversation for memory items.""" + + mode: Literal["phrase", "interval", "ai_determined"] + # For "phrase" mode: user says this to trigger memory check + trigger_phrase: Optional[str] = Field(default=None) + # For "interval" mode: check every N messages + message_interval: Optional[int] = Field(default=None, gt=0) + + # Note: Message counter is always per-session (in-memory), resets on agent restart diff --git a/src/vgm/config/VectorGraphConfig.py b/src/vgm/config/VectorGraphConfig.py new file mode 100644 index 0000000..2b28547 --- /dev/null +++ b/src/vgm/config/VectorGraphConfig.py @@ -0,0 +1,10 @@ +"""Vector and graph database configuration schema.""" + +from pydantic import BaseModel + + +class VectorGraphConfig(BaseModel): + """Configuration for vector and graph database connections.""" + + qdrant_collection: str = "vgm_memory" # Qdrant collection name + # JanusGraph doesn't need collection name, uses single graph with filtered queries diff --git a/src/vgm/config/__init__.py b/src/vgm/config/__init__.py new file mode 100644 index 0000000..d94e1df --- /dev/null +++ b/src/vgm/config/__init__.py @@ -0,0 +1,13 @@ +"""Configuration exports for Vector Graph Memory.""" + +from .MemoryConfig import MemoryConfig +from .MemoryTriggerConfig import MemoryTriggerConfig +from .AuditConfig import AuditConfig +from .VectorGraphConfig import VectorGraphConfig + +__all__ = [ + "MemoryConfig", + "MemoryTriggerConfig", + "AuditConfig", + "VectorGraphConfig", +] From dbaf62779ab925698700b5980ca597a7daffe040 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:21 -0700 Subject: [PATCH 15/21] feat: add audit logging system with multiple backends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement audit trail for memory operations: - Abstract AuditLogger interface - JSONL backend for file-based logging - MongoDB backend for centralized logging - Track all add/update/delete/search operations - Include metadata, affected entities, and user context 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- src/vgm/audit/AuditBackend.py | 30 ++++++++ src/vgm/audit/JSONLAuditBackend.py | 116 +++++++++++++++++++++++++++++ src/vgm/audit/MongoAuditBackend.py | 76 +++++++++++++++++++ src/vgm/audit/__init__.py | 11 +++ 4 files changed, 233 insertions(+) create mode 100644 src/vgm/audit/AuditBackend.py create mode 100644 src/vgm/audit/JSONLAuditBackend.py create mode 100644 src/vgm/audit/MongoAuditBackend.py create mode 100644 src/vgm/audit/__init__.py diff --git a/src/vgm/audit/AuditBackend.py b/src/vgm/audit/AuditBackend.py new file mode 100644 index 0000000..6428bae --- /dev/null +++ b/src/vgm/audit/AuditBackend.py @@ -0,0 +1,30 @@ +"""Abstract base class for audit backends.""" + +from abc import ABC, abstractmethod +from typing import List + +from ..schemas import AuditEntry + + +class AuditBackend(ABC): + """Abstract base class for audit log backends.""" + + @abstractmethod + def log_operation(self, entry: AuditEntry) -> None: + """Log a memory operation.""" + pass + + @abstractmethod + def get_recent(self, limit: int = 50) -> List[AuditEntry]: + """Get recent log entries.""" + pass + + @abstractmethod + def get_by_session(self, session_id: str) -> List[AuditEntry]: + """Get all logs for a specific session.""" + pass + + @abstractmethod + def get_entity_history(self, entity_id: str) -> List[AuditEntry]: + """Get all operations affecting a specific entity.""" + pass diff --git a/src/vgm/audit/JSONLAuditBackend.py b/src/vgm/audit/JSONLAuditBackend.py new file mode 100644 index 0000000..bcc5427 --- /dev/null +++ b/src/vgm/audit/JSONLAuditBackend.py @@ -0,0 +1,116 @@ +"""JSONL file-based audit backend.""" + +import os +from datetime import datetime, timezone +from pathlib import Path +from typing import List +from threading import Lock + +from pydantic import ValidationError + +from .AuditBackend import AuditBackend +from ..schemas import AuditEntry +from ..config import AuditConfig + + +class JSONLAuditBackend(AuditBackend): + """File-based JSONL audit logging with rotation support.""" + + def __init__(self, config: AuditConfig): + self.config = config + self.log_dir = Path(config.log_dir or os.path.expanduser("~/.vgm/logs")) + self.log_dir.mkdir(parents=True, exist_ok=True) + self.current_log_file = self.log_dir / "audit.jsonl" + self._write_lock = Lock() + + def _should_rotate(self) -> bool: + """Check if log file should be rotated.""" + if not self.current_log_file.exists(): + return False + + file_size_mb = self.current_log_file.stat().st_size / (1024 * 1024) + return file_size_mb >= self.config.rotation_size_mb + + def _rotate_log(self) -> None: + """Rotate log file based on configured period.""" + if not self.current_log_file.exists(): + return + + timestamp = datetime.now(timezone.utc) + if self.config.rotation_period == "daily": + suffix = timestamp.strftime("%Y-%m-%d") + elif self.config.rotation_period == "weekly": + suffix = timestamp.strftime("%Y-W%U") + else: # monthly + suffix = timestamp.strftime("%Y-%m") + + rotated_file = self.log_dir / f"audit-{suffix}.jsonl" + + # If rotated file exists, append to it; otherwise rename current + if rotated_file.exists(): + with open(rotated_file, "a") as dest, open(self.current_log_file, "r") as src: + dest.write(src.read()) + self.current_log_file.unlink() + else: + self.current_log_file.rename(rotated_file) + + def log_operation(self, entry: AuditEntry) -> None: + """Append audit entry to JSONL file.""" + with self._write_lock: + if self._should_rotate(): + self._rotate_log() + + with open(self.current_log_file, "a") as f: + json_line = entry.model_dump_json() + "\n" + f.write(json_line) + + def get_recent(self, limit: int = 50) -> List[AuditEntry]: + """Get most recent N entries.""" + if not self.current_log_file.exists(): + return [] + + entries = [] + with open(self.current_log_file, "r") as f: + lines = f.readlines() + # Take last N lines + for line in lines[-limit:]: + try: + entries.append(AuditEntry.model_validate_json(line)) + except (ValidationError, ValueError): + # Skip malformed lines - log parse errors are expected in corrupted files + continue + return entries + + def get_by_session(self, session_id: str) -> List[AuditEntry]: + """Get all entries for a session (scans entire file).""" + if not self.current_log_file.exists(): + return [] + + entries = [] + with open(self.current_log_file, "r") as f: + for line in f: + try: + entry = AuditEntry.model_validate_json(line) + if entry.session_id == session_id: + entries.append(entry) + except (ValidationError, ValueError): + # Skip malformed lines + continue + return entries + + def get_entity_history(self, entity_id: str) -> List[AuditEntry]: + """Get all operations affecting an entity (scans entire file).""" + if not self.current_log_file.exists(): + return [] + + entries = [] + with open(self.current_log_file, "r") as f: + for line in f: + try: + entry = AuditEntry.model_validate_json(line) + if entity_id in entry.affected_entities: + entries.append(entry) + except (ValidationError, ValueError): + # Skip malformed lines + continue + return entries diff --git a/src/vgm/audit/MongoAuditBackend.py b/src/vgm/audit/MongoAuditBackend.py new file mode 100644 index 0000000..24348be --- /dev/null +++ b/src/vgm/audit/MongoAuditBackend.py @@ -0,0 +1,76 @@ +"""MongoDB-based audit backend.""" + +from typing import List + +from .AuditBackend import AuditBackend +from ..schemas import AuditEntry +from ..config import AuditConfig + + +class MongoAuditBackend(AuditBackend): + """MongoDB-based audit logging with optional TTL.""" + + def __init__(self, config: AuditConfig): + try: + from pymongo import MongoClient, ASCENDING, IndexModel + from pymongo.errors import ConnectionFailure + except ImportError: + raise ImportError( + "MongoDB backend requires pymongo. Install with: pip install vector-graph-memory[mongodb]" + ) + + self.config = config + if not config.connection_string: + raise ValueError("MongoDB backend requires connection_string in AuditConfig") + + self.client = MongoClient(config.connection_string) + self.db = self.client[config.database] + self.collection = self.db[config.collection] + + # Create indexes + self._ensure_indexes() + + def _ensure_indexes(self) -> None: + """Create indexes for efficient querying.""" + from pymongo import ASCENDING, IndexModel + + indexes = [ + IndexModel([("timestamp", ASCENDING)]), + IndexModel([("session_id", ASCENDING)]), + IndexModel([("project_id", ASCENDING)]), + IndexModel([("affected_entities", ASCENDING)]), + ] + self.collection.create_indexes(indexes) + + # Create TTL index if configured + if self.config.ttl_days: + self.collection.create_index( + "timestamp", + expireAfterSeconds=self.config.ttl_days * 24 * 60 * 60 + ) + + def log_operation(self, entry: AuditEntry) -> None: + """Insert audit entry into MongoDB.""" + doc = entry.model_dump() + self.collection.insert_one(doc) + + def get_recent(self, limit: int = 50) -> List[AuditEntry]: + """Get most recent N entries.""" + docs = self.collection.find().sort("timestamp", -1).limit(limit) + return [AuditEntry.model_validate(doc) for doc in docs] + + def get_by_session(self, session_id: str) -> List[AuditEntry]: + """Get all entries for a session.""" + docs = self.collection.find({"session_id": session_id}).sort("timestamp", 1) + return [AuditEntry.model_validate(doc) for doc in docs] + + def get_entity_history(self, entity_id: str) -> List[AuditEntry]: + """Get all operations affecting an entity.""" + docs = self.collection.find( + {"affected_entities": entity_id} + ).sort("timestamp", 1) + return [AuditEntry.model_validate(doc) for doc in docs] + + def close(self) -> None: + """Close MongoDB connection.""" + self.client.close() diff --git a/src/vgm/audit/__init__.py b/src/vgm/audit/__init__.py new file mode 100644 index 0000000..43fd722 --- /dev/null +++ b/src/vgm/audit/__init__.py @@ -0,0 +1,11 @@ +"""Audit backend exports for Vector Graph Memory.""" + +from .AuditBackend import AuditBackend +from .JSONLAuditBackend import JSONLAuditBackend +from .MongoAuditBackend import MongoAuditBackend + +__all__ = [ + "AuditBackend", + "JSONLAuditBackend", + "MongoAuditBackend", +] From 9e6540197653a3980fd7a536e4b88d41d74964de Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:27 -0700 Subject: [PATCH 16/21] feat: add OpenAI-compatible REST API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implement FastAPI server with OpenAI compatibility: - /v1/chat/completions endpoint - /v1/models endpoint for model listing - Session-based agent management - Memory proposal confirmation endpoints - Environment-based configuration - CORS support for web clients Enables integration with Open WebUI and other OpenAI clients. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- src/vgm/api/__init__.py | 5 + src/vgm/api/server.py | 372 ++++++++++++++++++++++++++++++++++++++++ start_api.sh | 54 ++++++ 3 files changed, 431 insertions(+) create mode 100644 src/vgm/api/__init__.py create mode 100644 src/vgm/api/server.py create mode 100755 start_api.sh diff --git a/src/vgm/api/__init__.py b/src/vgm/api/__init__.py new file mode 100644 index 0000000..bd1e93c --- /dev/null +++ b/src/vgm/api/__init__.py @@ -0,0 +1,5 @@ +"""API module for Vector Graph Memory.""" + +from .server import app + +__all__ = ["app"] diff --git a/src/vgm/api/server.py b/src/vgm/api/server.py new file mode 100644 index 0000000..f76dad5 --- /dev/null +++ b/src/vgm/api/server.py @@ -0,0 +1,372 @@ +"""FastAPI server with OpenAI-compatible API for Open WebUI integration.""" + +import os +from typing import Dict, List, Optional, Literal, cast +from datetime import datetime +from contextlib import asynccontextmanager + +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel +from qdrant_client import QdrantClient +from gremlin_python.driver import client as gremlin_client +from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel +from dotenv import load_dotenv + +from ..MemoryAgent import MemoryAgent +from ..config import MemoryConfig, MemoryTriggerConfig, AuditConfig + + +# Load environment variables +load_dotenv() + + +# --- OpenAI-compatible API Models --- + +class ChatMessage(BaseModel): + """Chat message in OpenAI format.""" + role: str # system, user, assistant + content: str + + +class ChatCompletionRequest(BaseModel): + """OpenAI chat completion request.""" + model: str + messages: List[ChatMessage] + temperature: Optional[float] = 0.7 + max_tokens: Optional[int] = None + stream: Optional[bool] = False + user: Optional[str] = None # Used as session_id + + +class ChatCompletionChoice(BaseModel): + """Single completion choice.""" + index: int + message: ChatMessage + finish_reason: str + + +class ChatCompletionUsage(BaseModel): + """Token usage statistics.""" + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(BaseModel): + """OpenAI chat completion response.""" + id: str + object: str = "chat.completion" + created: int + model: str + choices: List[ChatCompletionChoice] + usage: ChatCompletionUsage + + +class ModelInfo(BaseModel): + """Model information.""" + id: str + object: str = "model" + created: int + owned_by: str = "vector-graph-memory" + + +class ModelList(BaseModel): + """List of available models.""" + object: str = "list" + data: List[ModelInfo] + + +# --- Global state --- + +class AppState: + """Application state container.""" + agent: Optional[MemoryAgent] = None + qdrant: Optional[QdrantClient] = None + janus: Optional[gremlin_client.Client] = None + session_proposals: Dict[str, List[str]] = {} # session_id -> list of proposal_ids + + +state = AppState() + + +# --- Lifecycle management --- + +@asynccontextmanager +async def lifespan(_app: FastAPI): + """Initialize and cleanup resources.""" + # Startup + import logging + + # Configure logging + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + ) + + print("🚀 Initializing Vector Graph Memory API...") + + # Connect to databases + qdrant_host = os.getenv('QDRANT_HOST', 'localhost') + qdrant_port = int(os.getenv('QDRANT_HTTP_PORT', '6333')) + janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost') + janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182')) + + state.qdrant = QdrantClient(host=qdrant_host, port=qdrant_port) + state.janus = gremlin_client.Client( + f'ws://{janusgraph_host}:{janusgraph_port}/gremlin', + 'g' + ) + + print(f" ✓ Connected to Qdrant at {qdrant_host}:{qdrant_port}") + print(f" ✓ Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}") + + # Initialize memory agent + embedding_model = OpenAIEmbeddingModel('text-embedding-3-small') + llm_model = os.getenv('LLM_MODEL', 'openai:gpt-4o-mini') + + memory_config = MemoryConfig( + use_case_description=os.getenv( + 'MEMORY_USE_CASE', + 'General purpose conversational memory' + ), + memory_threshold_description=os.getenv( + 'MEMORY_THRESHOLD', + 'Store important facts, decisions, and context from conversations' + ), + project_id=os.getenv('PROJECT_ID', 'default'), + similarity_threshold=float(os.getenv('SIMILARITY_THRESHOLD', '0.85')), + ) + + # Parse trigger mode with type safety + trigger_mode = os.getenv('TRIGGER_MODE', 'ai_determined') + valid_trigger_modes: tuple[Literal["phrase", "interval", "ai_determined"], ...] = ("phrase", "interval", "ai_determined") + if trigger_mode not in valid_trigger_modes: + trigger_mode = "ai_determined" + + # Parse trigger interval with safety for empty strings + trigger_interval_str = os.getenv('TRIGGER_INTERVAL', '0') + trigger_interval = int(trigger_interval_str) if trigger_interval_str else 0 + + trigger_config = MemoryTriggerConfig( + mode=cast(Literal["phrase", "interval", "ai_determined"], trigger_mode), + trigger_phrase=os.getenv('TRIGGER_PHRASE'), + message_interval=trigger_interval or None, + ) + + # Parse audit backend with type safety + audit_backend = os.getenv('AUDIT_BACKEND', 'jsonl') + valid_audit_backends: tuple[Literal["jsonl", "mongodb"], ...] = ("jsonl", "mongodb") + if audit_backend not in valid_audit_backends: + audit_backend = "jsonl" + + audit_config = AuditConfig( + backend=cast(Literal["jsonl", "mongodb"], audit_backend), + log_dir=os.getenv('AUDIT_LOG_DIR', '~/.vgm/logs'), + ) + + system_prompt = os.getenv( + 'SYSTEM_PROMPT', + 'You are a helpful AI assistant with persistent memory capabilities.' + ) + + state.agent = MemoryAgent( + qdrant_client=state.qdrant, + janus_client=state.janus, + embedding_model=embedding_model, + llm_model=llm_model, + system_prompt=system_prompt, + memory_config=memory_config, + trigger_config=trigger_config, + audit_config=audit_config, + ) + + print(f" ✓ Memory Agent initialized") + print(f" - Model: {llm_model}") + print(f" - Project: {memory_config.project_id}") + print(f" - Trigger: {trigger_config.mode}") + print(f"✓ API ready on http://0.0.0.0:8000") + + yield + + # Shutdown + print("\n🛑 Shutting down...") + if state.janus: + state.janus.close() + print("✓ Cleanup complete") + + +# --- FastAPI app --- + +app = FastAPI( + title="Vector Graph Memory API", + description="OpenAI-compatible API with persistent vector-graph memory", + version="1.0.0", + lifespan=lifespan, +) + + +# --- Endpoints --- + +@app.get("/") +async def root(): + """Health check endpoint.""" + return { + "status": "ok", + "service": "vector-graph-memory-api", + "version": "1.0.0", + } + + +@app.get("/v1/models") +async def list_models() -> ModelList: + """List available models (OpenAI compatible).""" + return ModelList( + data=[ + ModelInfo( + id="vector-graph-memory", + created=int(datetime.now().timestamp()), + ) + ] + ) + + +@app.post("/v1/chat/completions") +async def chat_completions(request: ChatCompletionRequest) -> ChatCompletionResponse: + """OpenAI-compatible chat completions endpoint.""" + if not state.agent: + raise HTTPException(status_code=503, detail="Agent not initialized") + + # Extract messages + if not request.messages: + raise HTTPException(status_code=400, detail="No messages provided") + + session_id = request.user or "default" + + # Build conversation context from message history + # Format: "Previous context:\nUser: ...\nAssistant: ...\n\nCurrent question: ..." + if len(request.messages) > 1: + # Multi-turn conversation - include history for context + context_parts = ["Here's our conversation so far:"] + for msg in request.messages[:-1]: + role = "User" if msg.role == "user" else "Assistant" + context_parts.append(f"{role}: {msg.content}") + context_parts.append(f"\nCurrent question: {request.messages[-1].content}") + user_message = "\n".join(context_parts) + else: + # Single message + user_message = request.messages[-1].content + + # Run agent + try: + result = state.agent.run(user_message, session_id=session_id) + assistant_response = result.output + + # Track any new proposals for this session + current_proposals = set(state.agent.pending_proposals.keys()) + session_proposals = set(state.session_proposals.get(session_id, [])) + new_proposals = current_proposals - session_proposals + + if new_proposals: + state.session_proposals[session_id] = list(current_proposals) + + except Exception as e: + raise HTTPException(status_code=500, detail=f"Agent error: {str(e)}") + + # Build response + return ChatCompletionResponse( + id=f"chatcmpl-{session_id}-{int(datetime.now().timestamp())}", + created=int(datetime.now().timestamp()), + model=request.model, + choices=[ + ChatCompletionChoice( + index=0, + message=ChatMessage(role="assistant", content=assistant_response), + finish_reason="stop", + ) + ], + usage=ChatCompletionUsage( + prompt_tokens=len(user_message.split()), # Rough estimate + completion_tokens=len(assistant_response.split()), + total_tokens=len(user_message.split()) + len(assistant_response.split()), + ), + ) + + +@app.post("/memory/confirm/{session_id}/{proposal_id}") +async def confirm_memory_proposal( + session_id: str, + proposal_id: str, + action: str = "add_new", + update_node_id: Optional[str] = None, +): + """Confirm a memory proposal. + + Args: + session_id: Session identifier + proposal_id: Proposal identifier from agent + action: 'add_new', 'update_existing', or 'cancel' + update_node_id: Node ID to update (if action='update_existing') + """ + if not state.agent: + raise HTTPException(status_code=503, detail="Agent not initialized") + + result = state.agent.confirm_memory_addition( + proposal_id=proposal_id, + action=action, + update_node_id=update_node_id, + ) + + # Update session proposals + if session_id in state.session_proposals: + state.session_proposals[session_id] = list(state.agent.pending_proposals.keys()) + + return {"status": "ok", "message": result} + + +@app.get("/memory/proposals/{session_id}") +async def get_pending_proposals(session_id: str): + """Get pending memory proposals for a session.""" + if not state.agent: + raise HTTPException(status_code=503, detail="Agent not initialized") + + session_proposal_ids = state.session_proposals.get(session_id, []) + proposals = {} + + for proposal_id in session_proposal_ids: + if proposal_id in state.agent.pending_proposals: + proposals[proposal_id] = state.agent.pending_proposals[proposal_id] + + return {"session_id": session_id, "proposals": proposals} + + +@app.get("/memory/audit/{session_id}") +async def get_session_audit(session_id: str, limit: int = 50): + """Get audit log for a session.""" + if not state.agent: + raise HTTPException(status_code=503, detail="Agent not initialized") + + entries = state.agent.get_audit_history(session_id=session_id, limit=limit) + + return { + "session_id": session_id, + "entries": [ + { + "timestamp": entry.timestamp.isoformat(), + "operation": entry.operation_type, + "summary": entry.summary, + "entities": entry.affected_entities, + } + for entry in entries + ], + } + + +if __name__ == "__main__": + import uvicorn + + uvicorn.run( + "server:app", + host="0.0.0.0", + port=8000, + reload=True, + ) diff --git a/start_api.sh b/start_api.sh new file mode 100755 index 0000000..5781030 --- /dev/null +++ b/start_api.sh @@ -0,0 +1,54 @@ +#!/bin/bash +# Start the Vector Graph Memory API server + +set -e + +echo "🚀 Starting Vector Graph Memory API..." + +# Check if .env exists +if [ ! -f .env ]; then + echo "⚠️ No .env file found. Copying from .env.example..." + cp .env.example .env + echo "⚠️ Please edit .env with your configuration and run again." + exit 1 +fi + +# Check if databases are running +echo "Checking database connections..." + +# Check Qdrant +QDRANT_HOST=${QDRANT_HOST:-localhost} +QDRANT_HTTP_PORT=${QDRANT_HTTP_PORT:-6333} + +if ! curl -s "http://$QDRANT_HOST:$QDRANT_HTTP_PORT/healthz" > /dev/null; then + echo "❌ Cannot connect to Qdrant at $QDRANT_HOST:$QDRANT_HTTP_PORT" + echo " Run: docker compose up -d" + exit 1 +fi +echo "✓ Qdrant is running" + +# Check JanusGraph +JANUSGRAPH_HOST=${JANUSGRAPH_HOST:-localhost} +JANUSGRAPH_PORT=${JANUSGRAPH_PORT:-8182} + +if ! timeout 2 bash -c "echo > /dev/tcp/$JANUSGRAPH_HOST/$JANUSGRAPH_PORT" 2>/dev/null; then + echo "❌ Cannot connect to JanusGraph at $JANUSGRAPH_HOST:$JANUSGRAPH_PORT" + echo " Run: docker compose up -d" + exit 1 +fi +echo "✓ JanusGraph is running" + +# Check if OPENAI_API_KEY is set +if [ -z "$OPENAI_API_KEY" ]; then + echo "❌ OPENAI_API_KEY not set in .env file" + exit 1 +fi +echo "✓ OpenAI API key configured" + +# Start the API server +echo "" +echo "Starting API server..." +python -m uvicorn src.vgm.api.server:app \ + --host ${API_HOST:-0.0.0.0} \ + --port ${API_PORT:-8000} \ + --reload From 903c31af9795ccc1b76b862360a75193d648c5d2 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:34 -0700 Subject: [PATCH 17/21] chore: add Docker containerization and deployment setup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add complete Docker stack: - Multi-stage Dockerfile for API service - docker-compose.yml with all services (Qdrant, JanusGraph, MongoDB, API, Open WebUI) - JanusGraph configuration and schema initialization - .env.example for environment configuration - Helper scripts for database initialization - .dockerignore for efficient builds Services are networked internally with only API and WebUI exposed. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- .dockerignore | 79 +++++++++++++++++++++++++ .env.example | 68 ++++++++++++++++++++++ Dockerfile | 30 ++++++++++ docker-compose.yml | 97 +++++++++++++++++++++++++++---- janusgraph-config.properties | 13 +++++ janusgraph-init.groovy | 48 +++++++++++++++ scripts/init_janusgraph_schema.py | 85 +++++++++++++++++++++++++++ 7 files changed, 409 insertions(+), 11 deletions(-) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 Dockerfile create mode 100644 janusgraph-config.properties create mode 100644 janusgraph-init.groovy create mode 100755 scripts/init_janusgraph_schema.py diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..be2119d --- /dev/null +++ b/.dockerignore @@ -0,0 +1,79 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# Virtual environments +venv/ +ENV/ +env/ +.venv + +# IDEs +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# Testing +.pytest_cache/ +.coverage +htmlcov/ +.tox/ + +# Jupyter +.ipynb_checkpoints/ +*.ipynb + +# Environment +.env +.env.local + +# Documentation +*.md +docs/ + +# Git +.git/ +.gitignore + +# Docker +Dockerfile +docker-compose.yml +.dockerignore + +# Logs and data +logs/ +*.log +.vgm/ + +# Database storage +qdrant_storage/ +janusgraph_storage/ +mongodb_storage/ + +# Scripts +start_api.sh +scripts/ + +# JanusGraph config +janusgraph-config.properties +janusgraph-init.groovy diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..e2420a4 --- /dev/null +++ b/.env.example @@ -0,0 +1,68 @@ +# Vector Graph Memory - Environment Configuration + +# === Database Configuration === +# NOTE: When running API via Docker, these are overridden in docker-compose.yml +# to use container names (qdrant, janusgraph) instead of localhost. +# These localhost values are for local development only. +QDRANT_HOST=localhost +QDRANT_HTTP_PORT=6333 +QDRANT_GRPC_PORT=6334 +JANUSGRAPH_HOST=localhost +JANUSGRAPH_PORT=8182 + +# === LLM Configuration === +# OpenAI API Key (required) +OPENAI_API_KEY=your_openai_api_key_here + +# LLM Model (supports PydanticAI format: openai:gpt-4o-mini, anthropic:claude-3-5-sonnet-20241022, etc.) +LLM_MODEL=openai:gpt-4o-mini + +# === Memory Configuration === +# Project identifier (unique namespace for memory storage) +PROJECT_ID=default + +# Use case description (guides the agent's memory behavior) +MEMORY_USE_CASE=General purpose conversational memory + +# Memory threshold description (what should be stored) +MEMORY_THRESHOLD=Store important facts, decisions, and context from conversations + +# Similarity threshold for duplicate detection (0.0-1.0) +SIMILARITY_THRESHOLD=0.85 + +# === Memory Trigger Configuration === +# Trigger mode: 'ai_determined', 'phrase', or 'interval' +TRIGGER_MODE=ai_determined + +# Trigger phrase (only used if TRIGGER_MODE=phrase) +# TRIGGER_PHRASE=save this to memory + +# Trigger interval in messages (only used if TRIGGER_MODE=interval) +# TRIGGER_INTERVAL=5 + +# === System Prompt === +SYSTEM_PROMPT=You are a helpful AI assistant with persistent memory capabilities. + +# === Audit Configuration === +# Audit backend: 'jsonl' or 'mongodb' +AUDIT_BACKEND=jsonl + +# Audit log directory (for JSONL backend) +AUDIT_LOG_DIR=~/.vgm/logs + +# MongoDB connection (only if AUDIT_BACKEND=mongodb) +# AUDIT_MONGODB_URI=mongodb://localhost:27017 +# AUDIT_MONGODB_DATABASE=vgm_audit +# AUDIT_MONGODB_COLLECTION=logs + +# === API Configuration === +# API server host and port +API_HOST=0.0.0.0 +API_PORT=8000 + +# === Docker Configuration === +# Qdrant version +QDRANT_VERSION=v1.12.5 + +# JanusGraph version +JANUSGRAPH_VERSION=latest diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..0f0a303 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,30 @@ +# Vector Graph Memory API Server +FROM python:3.11-slim + +# Set working directory +WORKDIR /app + +# Install system dependencies +RUN apt-get update && apt-get install -y \ + curl \ + && rm -rf /var/lib/apt/lists/* + +# Copy project files +COPY pyproject.toml . +COPY src/ src/ + +# Install Python dependencies +RUN pip install --no-cache-dir -e ".[api]" + +# Create directory for audit logs +RUN mkdir -p /app/logs + +# Expose API port +EXPOSE 8000 + +# Health check +HEALTHCHECK --interval=30s --timeout=10s --start-period=40s --retries=3 \ + CMD curl -f http://localhost:8000/ || exit 1 + +# Run the API server +CMD ["uvicorn", "src.vgm.api.server:app", "--host", "0.0.0.0", "--port", "8000"] diff --git a/docker-compose.yml b/docker-compose.yml index 2cd0cba..8a6abbc 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,12 +1,10 @@ -version: '3.8' - services: qdrant: image: qdrant/qdrant:${QDRANT_VERSION:-latest} container_name: ${QDRANT_CONTAINER_NAME:-vector-graph-qdrant} - ports: - - "${QDRANT_HTTP_PORT:-6333}:6333" # HTTP API - - "${QDRANT_GRPC_PORT:-6334}:6334" # gRPC API + expose: + - "6333" # HTTP API (internal only) + - "6334" # gRPC API (internal only) volumes: - qdrant_storage:/qdrant/storage environment: @@ -19,16 +17,89 @@ services: janusgraph: image: janusgraph/janusgraph:${JANUSGRAPH_VERSION:-latest} container_name: ${JANUSGRAPH_CONTAINER_NAME:-vector-graph-janusgraph} - ports: - - "${JANUSGRAPH_PORT:-8182}:8182" # Gremlin Server + expose: + - "8182" # Gremlin Server (internal only) volumes: - janusgraph_storage:/var/lib/janusgraph + - ./janusgraph-config.properties:/etc/opt/janusgraph/janusgraph.properties:ro + - ./janusgraph-init.groovy:/docker-entrypoint-initdb.d/init.groovy:ro environment: - JANUS_PROPS_TEMPLATE=berkeleyje-lucene - - janusgraph.storage.backend=berkeleyje - - janusgraph.storage.directory=/var/lib/janusgraph/data - - janusgraph.index.search.backend=lucene - - janusgraph.index.search.directory=/var/lib/janusgraph/index + restart: unless-stopped + networks: + - vector-graph-network + + mongodb: + image: mongo:${MONGODB_VERSION:-7} + container_name: ${MONGODB_CONTAINER_NAME:-vector-graph-mongodb} + expose: + - "27017" # MongoDB (internal only) + volumes: + - mongodb_storage:/data/db + environment: + - MONGO_INITDB_ROOT_USERNAME=${MONGODB_USER:-admin} + - MONGO_INITDB_ROOT_PASSWORD=${MONGODB_PASSWORD:-password} + restart: unless-stopped + networks: + - vector-graph-network + + api: + build: + context: . + dockerfile: Dockerfile + container_name: ${API_CONTAINER_NAME:-vector-graph-api} + ports: + - "${API_PORT:-8000}:8000" + volumes: + - ./logs:/app/logs + environment: + # Database connections (use container names) + - QDRANT_HOST=qdrant + - QDRANT_HTTP_PORT=6333 + - JANUSGRAPH_HOST=janusgraph + - JANUSGRAPH_PORT=8182 + + # LLM Configuration + - OPENAI_API_KEY=${OPENAI_API_KEY} + - LLM_MODEL=${LLM_MODEL:-openai:gpt-4o-mini} + + # Memory Configuration + - PROJECT_ID=${PROJECT_ID:-default} + - MEMORY_USE_CASE=${MEMORY_USE_CASE:-General purpose conversational memory} + - MEMORY_THRESHOLD=${MEMORY_THRESHOLD:-Store important facts, decisions, and context from conversations} + - SIMILARITY_THRESHOLD=${SIMILARITY_THRESHOLD:-0.85} + + # Trigger Configuration + - TRIGGER_MODE=${TRIGGER_MODE:-ai_determined} + - TRIGGER_PHRASE=${TRIGGER_PHRASE} + - TRIGGER_INTERVAL=${TRIGGER_INTERVAL} + + # System Prompt + - SYSTEM_PROMPT=${SYSTEM_PROMPT:-You are a helpful AI assistant with persistent memory capabilities.} + + # Audit Configuration + - AUDIT_BACKEND=${AUDIT_BACKEND:-jsonl} + - AUDIT_LOG_DIR=/app/logs + depends_on: + - qdrant + - janusgraph + restart: unless-stopped + networks: + - vector-graph-network + + open-webui: + image: ghcr.io/open-webui/open-webui:main + container_name: ${WEBUI_CONTAINER_NAME:-vector-graph-webui} + ports: + - "${WEBUI_PORT:-3000}:8080" + volumes: + - open_webui_data:/app/backend/data + environment: + - OPENAI_API_BASE_URLS=http://api:8000/v1 + - OPENAI_API_KEYS=${OPENAI_API_KEY} + - WEBUI_AUTH=false + depends_on: + - api restart: unless-stopped networks: - vector-graph-network @@ -38,6 +109,10 @@ volumes: driver: local janusgraph_storage: driver: local + mongodb_storage: + driver: local + open_webui_data: + driver: local networks: vector-graph-network: diff --git a/janusgraph-config.properties b/janusgraph-config.properties new file mode 100644 index 0000000..5bf5b6d --- /dev/null +++ b/janusgraph-config.properties @@ -0,0 +1,13 @@ +# JanusGraph configuration with schema constraints disabled +gremlin.graph=org.janusgraph.core.JanusGraphFactory + +# Storage backend +storage.backend=berkeleyje +storage.directory=/var/lib/janusgraph/data + +# Index backend +index.search.backend=lucene +index.search.directory=/var/lib/janusgraph/index + +# Disable schema constraints - allow dynamic vertex/edge labels +schema.default=default diff --git a/janusgraph-init.groovy b/janusgraph-init.groovy new file mode 100644 index 0000000..f1d6bbe --- /dev/null +++ b/janusgraph-init.groovy @@ -0,0 +1,48 @@ +// JanusGraph schema initialization script +// This runs automatically when JanusGraph starts + +mgmt = graph.openManagement() + +// Create vertex label 'node' if it doesn't exist +if (mgmt.getVertexLabel('node') == null) { + mgmt.makeVertexLabel('node').make() +} + +// Create edge label 'relationship' if it doesn't exist +if (mgmt.getEdgeLabel('relationship') == null) { + mgmt.makeEdgeLabel('relationship').make() +} + +// Create property keys +if (mgmt.getPropertyKey('node_id') == null) { + mgmt.makePropertyKey('node_id').dataType(String.class).make() +} +if (mgmt.getPropertyKey('node_type') == null) { + mgmt.makePropertyKey('node_type').dataType(String.class).make() +} +if (mgmt.getPropertyKey('project_id') == null) { + mgmt.makePropertyKey('project_id').dataType(String.class).make() +} +if (mgmt.getPropertyKey('created_at') == null) { + mgmt.makePropertyKey('created_at').dataType(String.class).make() +} +if (mgmt.getPropertyKey('updated_at') == null) { + mgmt.makePropertyKey('updated_at').dataType(String.class).make() +} +if (mgmt.getPropertyKey('edge_id') == null) { + mgmt.makePropertyKey('edge_id').dataType(String.class).make() +} +if (mgmt.getPropertyKey('relationship_type') == null) { + mgmt.makePropertyKey('relationship_type').dataType(String.class).make() +} +if (mgmt.getPropertyKey('description') == null) { + mgmt.makePropertyKey('description').dataType(String.class).make() +} +if (mgmt.getPropertyKey('source') == null) { + mgmt.makePropertyKey('source').dataType(String.class).make() +} +if (mgmt.getPropertyKey('confidence') == null) { + mgmt.makePropertyKey('confidence').dataType(Double.class).make() +} + +mgmt.commit() diff --git a/scripts/init_janusgraph_schema.py b/scripts/init_janusgraph_schema.py new file mode 100755 index 0000000..34bca05 --- /dev/null +++ b/scripts/init_janusgraph_schema.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Initialize JanusGraph schema. + +Run this script once before using the library to create the required +vertex labels, edge labels, and property keys in JanusGraph. + +Usage: + python scripts/init_janusgraph_schema.py +""" + +import os +import sys +from pathlib import Path + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent.parent / "src")) + +from gremlin_python.driver import client as gremlin_client + +def init_schema(): + """Initialize JanusGraph schema.""" + # Connect to JanusGraph + janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost') + janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182')) + + print(f"Connecting to JanusGraph at {janusgraph_host}:{janusgraph_port}...") + janus = gremlin_client.Client( + f'ws://{janusgraph_host}:{janusgraph_port}/gremlin', + 'g' + ) + + schema_script = """ +mgmt = graph.openManagement() +if (mgmt.getVertexLabel('node') == null) { + mgmt.makeVertexLabel('node').make() +} +if (mgmt.getEdgeLabel('relationship') == null) { + mgmt.makeEdgeLabel('relationship').make() +} +if (mgmt.getPropertyKey('node_id') == null) { + mgmt.makePropertyKey('node_id').dataType(String.class).make() +} +if (mgmt.getPropertyKey('node_type') == null) { + mgmt.makePropertyKey('node_type').dataType(String.class).make() +} +if (mgmt.getPropertyKey('project_id') == null) { + mgmt.makePropertyKey('project_id').dataType(String.class).make() +} +if (mgmt.getPropertyKey('created_at') == null) { + mgmt.makePropertyKey('created_at').dataType(String.class).make() +} +if (mgmt.getPropertyKey('updated_at') == null) { + mgmt.makePropertyKey('updated_at').dataType(String.class).make() +} +if (mgmt.getPropertyKey('edge_id') == null) { + mgmt.makePropertyKey('edge_id').dataType(String.class).make() +} +if (mgmt.getPropertyKey('relationship_type') == null) { + mgmt.makePropertyKey('relationship_type').dataType(String.class).make() +} +if (mgmt.getPropertyKey('description') == null) { + mgmt.makePropertyKey('description').dataType(String.class).make() +} +if (mgmt.getPropertyKey('source') == null) { + mgmt.makePropertyKey('source').dataType(String.class).make() +} +if (mgmt.getPropertyKey('confidence') == null) { + mgmt.makePropertyKey('confidence').dataType(Double.class).make() +} +mgmt.commit() +'Schema initialized successfully' +""" + + print("Creating JanusGraph schema...") + try: + result = janus.submit(schema_script).all().result() + print(f"✓ {result[0]}") + except Exception as e: + print(f"✗ Error: {e}") + sys.exit(1) + finally: + janus.close() + +if __name__ == "__main__": + init_schema() From b25bf6108b15b1b8f23f5430ec662d0df2a3cf71 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:40 -0700 Subject: [PATCH 18/21] chore: update project dependencies and build configuration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Update pyproject.toml: - Add all required dependencies (pydantic, pymongo, nest-asyncio) - Add optional dependencies for MongoDB and API - Configure hatchling build system for src-layout - Update project description 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- pyproject.toml | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a848a89..bffdc72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,14 +1,27 @@ [project] name = "vector-graph-memory" version = "0.1.0" -description = "Add your description here" +description = "Hybrid vector-graph database backend for AI agents with persistent memory" requires-python = ">=3.11" dependencies = [ "gremlinpython>=3.8.0", + "nest-asyncio>=1.5.0", + "pydantic>=2.0.0", "pydantic-ai>=1.70.0", + "pymongo>=4.16.0", "qdrant-client>=1.17.1", ] +[project.optional-dependencies] +mongodb = [ + "pymongo>=4.6.0", +] +api = [ + "fastapi>=0.115.0", + "uvicorn[standard]>=0.32.0", + "python-dotenv>=1.0.0", +] + [dependency-groups] dev = [ "ipykernel>=7.2.0", @@ -17,3 +30,10 @@ dev = [ "pytest>=9.0.2", "pytest-cov>=7.1.0", ] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/vgm"] From bb8c31fc2f996a316c5a48d192d1c55acf007593 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Sat, 21 Mar 2026 19:59:46 -0700 Subject: [PATCH 19/21] docs: update README and add API documentation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Major documentation updates: - README.md: Add API usage, Open WebUI integration, status updates - API.md: Complete OpenAI-compatible API documentation - playground.ipynb: Updated with MemoryAgent examples - playground_api.ipynb: New notebook for REST API examples Documents the complete v1.0 feature set and integration options. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- API.md | 418 ++++++++++++++++++++++++++++++ README.md | 101 ++++++-- playground.ipynb | 521 +++++++++++++++++++++++++++++++------ playground_api.ipynb | 595 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 1540 insertions(+), 95 deletions(-) create mode 100644 API.md create mode 100644 playground_api.ipynb diff --git a/API.md b/API.md new file mode 100644 index 0000000..f9e402e --- /dev/null +++ b/API.md @@ -0,0 +1,418 @@ +# Vector Graph Memory API + +OpenAI-compatible API server for integrating Vector Graph Memory with Open WebUI and other LLM frontends. + +## Quick Start + +### Option 1: Docker (Recommended) + +Run everything with Docker Compose: + +```bash +# 1. Configure environment +cp .env.example .env +# Edit .env - at minimum set OPENAI_API_KEY + +# 2. Start all services (databases + API) +docker compose up -d + +# 3. Check logs +docker compose logs -f api +``` + +The API will be available at `http://localhost:8000` + +**Included services:** +- Qdrant (vector database) +- JanusGraph (graph database) +- MongoDB (optional, for audit logs) +- API server (FastAPI with memory agent) + +**How it works:** +- All containers run on an internal Docker network (`vector-graph-network`) +- The API container connects to databases using container names (`qdrant`, `janusgraph`) +- Only the API port (8000) is exposed to your host machine +- Database connections are handled automatically - no localhost configuration needed! + +### Option 2: Local Development + +Run the API locally while databases run in Docker: + +```bash +# 1. Install dependencies +pip install -e ".[api]" + +# 2. Configure environment +cp .env.example .env +# Edit .env with your settings + +# 3. Start only databases +docker compose up -d qdrant janusgraph + +# 4. Start API locally +./start_api.sh +``` + +The API will be available at `http://localhost:8000` + +## Open WebUI Integration + +### Add as External Connection + +1. Open your Open WebUI instance +2. Go to **Settings** → **Connections** +3. Add a new **OpenAI API** connection: + - **Name**: Vector Graph Memory + - **Base URL**: `http://localhost:8000/v1` + - **API Key**: (any value, not validated) +4. Save and select the "vector-graph-memory" model + +### Usage + +Simply chat with the agent through Open WebUI. The agent will: + +- Automatically search memory for relevant context +- Propose storing important information +- Track conversations in audit logs + +**Memory Trigger Modes**: + +- `ai_determined` (default): Agent decides when to check memory +- `phrase`: Trigger on specific phrase (e.g., "save this to memory") +- `interval`: Check every N messages + +Configure via `TRIGGER_MODE` in `.env` + +## API Endpoints + +### OpenAI-Compatible Endpoints + +#### `POST /v1/chat/completions` + +Standard OpenAI chat completions endpoint. + +**Request:** +```json +{ + "model": "vector-graph-memory", + "messages": [ + {"role": "user", "content": "What do you remember about my job search?"} + ], + "user": "optional-session-id" +} +``` + +**Response:** +```json +{ + "id": "chatcmpl-...", + "object": "chat.completion", + "created": 1234567890, + "model": "vector-graph-memory", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Based on my memory, you applied to..." + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30 + } +} +``` + +#### `GET /v1/models` + +List available models. + +**Response:** +```json +{ + "object": "list", + "data": [ + { + "id": "vector-graph-memory", + "object": "model", + "created": 1234567890, + "owned_by": "vector-graph-memory" + } + ] +} +``` + +### Memory Management Endpoints + +#### `GET /memory/proposals/{session_id}` + +Get pending memory proposals for a session. + +**Response:** +```json +{ + "session_id": "user-123", + "proposals": { + "proposal-uuid-1": { + "content": "Applied to Google for Senior SWE role", + "entity_type": "job", + "relationships": [], + "similar_nodes": [] + } + } +} +``` + +#### `POST /memory/confirm/{session_id}/{proposal_id}` + +Confirm or reject a memory proposal. + +**Query Parameters:** +- `action`: `add_new`, `update_existing`, or `cancel` +- `update_node_id`: Required if action is `update_existing` + +**Response:** +```json +{ + "status": "ok", + "message": "Successfully added job to memory with ID: abc-123" +} +``` + +#### `GET /memory/audit/{session_id}` + +Get audit log for a session. + +**Query Parameters:** +- `limit`: Maximum number of entries (default: 50) + +**Response:** +```json +{ + "session_id": "user-123", + "entries": [ + { + "timestamp": "2026-03-22T01:55:04", + "operation": "add_node", + "summary": "Added job: Applied to Google...", + "entities": ["node-uuid-1"] + } + ] +} +``` + +### Health Check + +#### `GET /` + +API health check. + +**Response:** +```json +{ + "status": "ok", + "service": "vector-graph-memory-api", + "version": "1.0.0" +} +``` + +## Configuration Reference + +### Environment Variables + +See `.env.example` for full configuration options. + +**Key Settings:** + +| Variable | Description | Default | +|----------|-------------|---------| +| `LLM_MODEL` | PydanticAI model string | `openai:gpt-4o-mini` | +| `PROJECT_ID` | Memory namespace | `default` | +| `MEMORY_USE_CASE` | Use case description | General purpose memory | +| `TRIGGER_MODE` | When to check memory | `ai_determined` | +| `SIMILARITY_THRESHOLD` | Duplicate detection threshold | `0.85` | + +### Memory Trigger Modes + +1. **AI Determined** (`ai_determined`) + - Agent autonomously decides when to check memory + - Best for general use cases + - Set: `TRIGGER_MODE=ai_determined` + +2. **Phrase-based** (`phrase`) + - Trigger on specific phrase + - User must explicitly request memory storage + - Set: `TRIGGER_MODE=phrase` and `TRIGGER_PHRASE=save this to memory` + +3. **Interval-based** (`interval`) + - Check every N messages + - Predictable behavior + - Set: `TRIGGER_MODE=interval` and `TRIGGER_INTERVAL=5` + +## Architecture + +The API server: + +1. **Initializes on startup**: + - Connects to Qdrant and JanusGraph + - Creates MemoryAgent instance + - Loads configuration from environment + +2. **Handles requests**: + - Receives chat messages via OpenAI API + - Runs agent with memory tools + - Returns responses + +3. **Manages sessions**: + - Uses `user` field as session ID + - Tracks pending proposals per session + - Maintains conversation context + +4. **Provides memory control**: + - Endpoints to view/confirm proposals + - Audit log access + - Session management + +## Development + +### Running with Auto-reload + +```bash +python -m uvicorn src.vgm.api.server:app --reload --host 0.0.0.0 --port 8000 +``` + +### Testing the API + +```bash +# Health check +curl http://localhost:8000/ + +# Chat completion +curl -X POST http://localhost:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vector-graph-memory", + "messages": [{"role": "user", "content": "Hello!"}], + "user": "test-session" + }' + +# Check proposals +curl http://localhost:8000/memory/proposals/test-session +``` + +### API Documentation + +Interactive API docs available at: +- Swagger UI: http://localhost:8000/docs +- ReDoc: http://localhost:8000/redoc + +## Docker Commands + +### Start all services +```bash +docker compose up -d +``` + +### View logs +```bash +# All services +docker compose logs -f + +# Just API +docker compose logs -f api + +# Just databases +docker compose logs -f qdrant janusgraph +``` + +### Rebuild API after code changes +```bash +docker compose up -d --build api +``` + +### Stop all services +```bash +docker compose down +``` + +### Stop and remove data +```bash +docker compose down -v +``` + +### Restart just the API +```bash +docker compose restart api +``` + +## Troubleshooting + +### Cannot connect to databases + +Error: `Cannot connect to Qdrant/JanusGraph` + +**Solution:** +```bash +docker compose up -d +# Wait 10-15 seconds for JanusGraph to initialize + +# Check if services are healthy +docker compose ps +``` + +### API key not set + +Error: `OPENAI_API_KEY not set` + +**Solution:** +Add to `.env`: +```bash +OPENAI_API_KEY=sk-... +``` + +Then restart: +```bash +docker compose restart api +``` + +### Port already in use + +Error: `Address already in use` + +**Solution:** +Change port in `.env`: +```bash +API_PORT=8001 +``` + +Then rebuild: +```bash +docker compose up -d +``` + +### API container won't start + +**Check logs:** +```bash +docker compose logs api +``` + +**Common issues:** +- Missing `OPENAI_API_KEY` in `.env` +- Database services not ready (wait 15-20 seconds) +- Port conflict (change `API_PORT`) + +**Force rebuild:** +```bash +docker compose up -d --build --force-recreate api +``` + +## Next Steps + +- See `playground.ipynb` for usage examples +- Check `README.md` for overall project documentation +- Review `.env.example` for all configuration options diff --git a/README.md b/README.md index e18ce8b..43a259c 100644 --- a/README.md +++ b/README.md @@ -12,11 +12,12 @@ A hybrid vector-graph database backend for AI agents, providing long-term persis ## Architecture -- **Type:** Python library (v1), with potential service wrapper in future versions -- **Python Version:** 3.14+ +- **Type:** Python library + OpenAI-compatible REST API +- **Python Version:** 3.11+ - **Vector Database:** Qdrant (semantic entry point discovery) - **Graph Database:** JanusGraph (relationship storage and traversal) - **AI Framework:** PydanticAI +- **API Framework:** FastAPI with OpenAI-compatible endpoints ## Core Concepts @@ -66,43 +67,109 @@ The system leverages the strengths of both database paradigms: - Python 3.14+ - uv (Python package manager) -### Database Setup +### Quick Start -Start the required databases using Docker Compose: +**Option 1: Using the API (Recommended)** ```bash +# 1. Configure environment +cp .env.example .env +# Edit .env - at minimum set OPENAI_API_KEY + +# 2. Start all services (databases + API) docker compose up -d + +# 3. Try the API +curl http://localhost:8000/v1/models ``` -This will start: -- **Qdrant** on ports 6333 (HTTP) and 6334 (gRPC) -- **JanusGraph** on port 8182 (Gremlin Server) +The API will be available at `http://localhost:8000` with OpenAI-compatible endpoints. -To stop the databases: +**Option 2: Using the Python Library** ```bash -docker compose down +# 1. Start databases only +docker compose up -d qdrant janusgraph + +# 2. Install library +pip install -e . + +# 3. Use in Python +# See playground.ipynb for examples ``` -To remove all data volumes: +### Playground Notebooks + +Two Jupyter notebooks are provided: + +- **`playground_api.ipynb`** - Examples using the REST API (recommended) +- **`playground.ipynb`** - Examples using the Python library directly + +Start Jupyter and open either notebook: ```bash -docker compose down -v +jupyter notebook ``` -### Environment Configuration +## Integration with Open WebUI + +The Vector Graph Memory API is compatible with Open WebUI and other OpenAI-compatible clients. + +### Using the Included Open WebUI Instance -Copy the example environment file and configure as needed: +The easiest way to test the API is with the included Open WebUI container: ```bash -cp .env.example .env +# Start all services including Open WebUI +docker compose up -d + +# Open WebUI will be available at http://localhost:3000 ``` -## Status +The Open WebUI instance is pre-configured to use the Vector Graph Memory API. Simply: + +1. Navigate to `http://localhost:3000` in your browser +2. Start chatting - the agent has persistent memory capabilities +3. The agent will propose adding information to memory during conversations +4. Confirm memory proposals through the API endpoints (see `playground_api.ipynb` for examples) + +### Using an External Open WebUI Instance + +If you have your own Open WebUI installation: + +1. Start the API: `docker compose up -d api` +2. In Open WebUI, add a new OpenAI API connection: + - **Base URL:** `http://localhost:8000/v1` (or `http://api:8000/v1` if on the same Docker network) + - **API Key:** (any value - authentication is disabled by default) + - **Model:** `vector-graph-memory` + +### Configuration -**Current Phase:** Planning and design +Customize Open WebUI port in `.env`: + +```bash +WEBUI_PORT=3000 # Default port +``` + +See [API.md](API.md) for complete API documentation. + +## Status -This project is in active development. The README will be updated as implementation progresses. +**Current Version:** v1.0 - Core functionality complete + +**Features:** +- ✅ Hybrid vector-graph storage +- ✅ PydanticAI agent with memory tools +- ✅ OpenAI-compatible REST API +- ✅ Docker containerization +- ✅ Memory proposals and confirmations +- ✅ Audit logging (JSONL/MongoDB) +- ✅ Configurable memory triggers + +**In Development:** +- Temporal tracking and recency scoring +- Memory consolidation +- Advanced graph traversal patterns ## License diff --git a/playground.ipynb b/playground.ipynb index 83245e8..059fc0f 100644 --- a/playground.ipynb +++ b/playground.ipynb @@ -2,173 +2,538 @@ "cells": [ { "cell_type": "markdown", + "id": "header", "metadata": {}, "source": [ - "# Vector Graph Memory Playground\n", + "# Vector Graph Memory - Playground\n", "\n", - "This notebook demonstrates the basic setup of a PydanticAI agent with Qdrant and JanusGraph database tools." - ], - "id": "959384cb2a852b81" + "This notebook demonstrates the Vector Graph Memory system with a job search use case." + ] }, { "cell_type": "markdown", + "id": "setup-header", "metadata": {}, "source": [ "## Setup and Imports" - ], - "id": "3f1ac3ddc5cc228b" + ] }, { "cell_type": "code", + "id": "imports", "metadata": { "ExecuteTime": { - "end_time": "2026-03-21T23:12:25.707261937Z", - "start_time": "2026-03-21T23:12:23.803285021Z" + "end_time": "2026-03-22T02:06:41.351851630Z", + "start_time": "2026-03-22T02:06:41.295520936Z" } }, - "source": "import os\nfrom dotenv import load_dotenv\nfrom qdrant_client import QdrantClient\nfrom gremlin_python.driver import client as gremlin_client\nfrom pydantic_ai import Agent, RunContext\nfrom pydantic_ai.models import infer_model\nfrom pydantic import BaseModel\n\nfrom vector_graph_memory.memory import VectorGraphMemory\n\n# Load environment variables\nload_dotenv()", - "id": "dd65b95cb0cd2644", - "outputs": [], - "execution_count": null + "source": "import os\nfrom dotenv import load_dotenv\nfrom qdrant_client import QdrantClient\nfrom gremlin_python.driver import client as gremlin_client\nfrom pydantic_ai.embeddings.openai import OpenAIEmbeddingModel\n\nfrom src.vgm import MemoryAgent, MemoryConfig, MemoryTriggerConfig, AuditConfig\n\n# Load environment variables\nload_dotenv()\n\nprint(\"✓ Imports successful\")", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Imports successful\n" + ] + } + ], + "execution_count": 12 }, { "cell_type": "markdown", + "id": "db-header", "metadata": {}, "source": [ - "## Database Connections" - ], - "id": "1d9228511ff2cb25" + "## Database Connections\n", + "\n", + "Make sure you've started the databases with `docker compose up -d`" + ] }, { "cell_type": "code", - "metadata": {}, + "id": "qdrant", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:42.403795971Z", + "start_time": "2026-03-22T02:06:42.332203557Z" + } + }, "source": [ "# Qdrant connection\n", "qdrant_host = os.getenv('QDRANT_HOST', 'localhost')\n", - "qdrant_port = int(os.getenv('QDRANT_PORT', '6333'))\n", + "qdrant_port = int(os.getenv('QDRANT_HTTP_PORT', '6333'))\n", "\n", "qdrant = QdrantClient(host=qdrant_host, port=qdrant_port)\n", - "print(f\"Connected to Qdrant at {qdrant_host}:{qdrant_port}\")" + "print(f\"✓ Connected to Qdrant at {qdrant_host}:{qdrant_port}\")" ], - "id": "e9dc7a4addca6ad5", - "outputs": [], - "execution_count": null + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Connected to Qdrant at localhost:6123\n" + ] + } + ], + "execution_count": 13 }, { "cell_type": "code", - "metadata": {}, - "source": "# JanusGraph connection\njanusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\njanusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n\njanus = gremlin_client.Client(\n f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n 'g'\n)\nprint(f\"Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")", - "id": "173bac3780209c1d", - "outputs": [], - "execution_count": null + "id": "janusgraph", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:42.926953548Z", + "start_time": "2026-03-22T02:06:42.869970038Z" + } + }, + "source": [ + "# JanusGraph connection\n", + "janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost')\n", + "janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182'))\n", + "\n", + "janus = gremlin_client.Client(\n", + " f'ws://{janusgraph_host}:{janusgraph_port}/gremlin',\n", + " 'g'\n", + ")\n", + "print(f\"✓ Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Connected to JanusGraph at localhost:8481\n" + ] + } + ], + "execution_count": 14 }, { "cell_type": "markdown", + "id": "agent-header", "metadata": {}, "source": [ - "## Database Context" - ], - "id": "15d850aa831b42e1" + "## Initialize Memory Agent\n", + "\n", + "Configure the agent with:\n", + "- Embedding model for vector search\n", + "- LLM model for agent reasoning\n", + "- Memory configuration (use case, threshold, project)\n", + "- Audit logging (JSONL or MongoDB)" + ] }, { "cell_type": "code", + "id": "agent-init", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:44.946009691Z", + "start_time": "2026-03-22T02:06:44.254525980Z" + } + }, + "source": "# Configure embedding and LLM models\nembedding_model = OpenAIEmbeddingModel('text-embedding-3-small')\nllm_model = 'openai:gpt-4o-mini'\n\n# Define system prompt\nsystem_prompt = \"\"\"\nYou are a helpful assistant for managing job search activities.\nYou help track job applications, companies, contacts, and interactions.\n\"\"\"\n\n# Configure memory behavior\nmemory_config = MemoryConfig(\n use_case_description=\"Track job search: companies, job postings, contacts, and interactions\",\n memory_threshold_description=\"Store specific job postings, company information, contact details, and meaningful interactions\",\n project_id=\"job_search_2026\",\n similarity_threshold=0.85,\n)\n\n# Configure memory triggers (when to check for items to remember)\ntrigger_config = MemoryTriggerConfig(\n mode=\"phrase\",\n trigger_phrase=\"save this to memory\",\n)\n\n# Configure audit logging (using JSONL by default)\naudit_config = AuditConfig(\n backend=\"jsonl\",\n log_dir=\"~/.vgm/logs\",\n)\n\n# Create the agent\nagent = MemoryAgent(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_model=embedding_model,\n llm_model=llm_model,\n system_prompt=system_prompt,\n memory_config=memory_config,\n trigger_config=trigger_config,\n audit_config=audit_config,\n)\n\nprint(\"✓ Memory Agent initialized\")\nprint(f\" - Embedding: text-embedding-3-small\")\nprint(f\" - LLM: {llm_model}\")\nprint(f\" - Project: {memory_config.project_id}\")\nprint(f\" - Trigger: {trigger_config.mode}\")\nprint(f\" - Audit: {audit_config.backend}\")", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Memory Agent initialized\n", + " - Embedding: text-embedding-3-small\n", + " - LLM: openai:gpt-4o-mini\n", + " - Project: job_search_2026\n", + " - Trigger: phrase\n", + " - Audit: jsonl\n" + ] + } + ], + "execution_count": 15 + }, + { + "cell_type": "markdown", + "id": "usage-header", "metadata": {}, - "source": "class DatabaseContext(BaseModel):\n \"\"\"Context containing memory system for the agent.\"\"\"\n memory: VectorGraphMemory\n \n class Config:\n arbitrary_types_allowed = True", - "id": "3471e89c4036c88d", - "outputs": [], - "execution_count": null + "source": [ + "## Usage Examples" + ] }, { "cell_type": "markdown", + "id": "ex1-header", "metadata": {}, "source": [ - "## Agent Tools" - ], - "id": "2fdd05d5a60b885c" + "### Example 1: Basic Conversation" + ] }, { "cell_type": "code", - "metadata": {}, - "source": "# Create the agent - choose your LLM model\n# Examples: 'openai:gpt-4', 'anthropic:claude-3-5-sonnet-20241022', 'ollama:llama3'\nagent = Agent(\n 'openai:gpt-4',\n deps_type=DatabaseContext,\n system_prompt=\"You are a helpful assistant with access to a vector-graph memory system.\"\n)\n\n@agent.tool\nasync def search_vector_memory(ctx: RunContext[DatabaseContext], query: str, limit: int = 5) -> str:\n \"\"\"\n Search the vector database for semantically similar content.\n \n Args:\n query: The search query text\n limit: Maximum number of results to return\n \n Returns:\n Search results as a formatted string\n \"\"\"\n results = await ctx.deps.memory.search_similar(query, limit=limit)\n \n if not results:\n return f\"No results found for query: '{query}'\"\n \n formatted_results = []\n for i, result in enumerate(results, 1):\n formatted_results.append(\n f\"{i}. [{result['type']}] {result['content']} (score: {result['score']:.3f})\"\n )\n \n return \"\\\\n\".join(formatted_results)\n\n@agent.tool\nasync def query_graph(ctx: RunContext[DatabaseContext], start_id: str, traversal_steps: str) -> str:\n \"\"\"\n Execute a graph traversal starting from a specific entity.\n \n Args:\n start_id: Starting entity ID\n traversal_steps: Gremlin traversal steps (e.g., \"out('works_at').values('content')\")\n \n Returns:\n Query results as a formatted string\n \"\"\"\n try:\n results = await ctx.deps.memory.query_graph_traversal(start_id, traversal_steps)\n return f\"Graph traversal results: {results}\"\n except Exception as e:\n return f\"Error executing graph traversal: {str(e)}\"\n\n@agent.tool\nasync def add_to_memory(ctx: RunContext[DatabaseContext], content: str, entity_type: str, metadata: dict | None = None) -> str:\n \"\"\"\n Add new content to both vector and graph databases.\n \n Args:\n content: The content to store\n entity_type: The type of entity (e.g., 'job', 'company', 'person')\n metadata: Optional additional metadata\n \n Returns:\n Confirmation message with entity ID\n \"\"\"\n entity_id = await ctx.deps.memory.add_entity(content, entity_type, metadata)\n return f\"Added {entity_type} to memory with ID: {entity_id}\"\n\n@agent.tool\nasync def link_entities(ctx: RunContext[DatabaseContext], from_id: str, to_id: str, relationship: str) -> str:\n \"\"\"\n Create a relationship between two entities in the graph.\n \n Args:\n from_id: Source entity ID\n to_id: Target entity ID\n relationship: Relationship type (e.g., 'works_at', 'applied_to')\n \n Returns:\n Confirmation message\n \"\"\"\n try:\n await ctx.deps.memory.add_relationship(from_id, to_id, relationship)\n return f\"Created '{relationship}' relationship from {from_id} to {to_id}\"\n except Exception as e:\n return f\"Error creating relationship: {str(e)}\"\n\nprint(\"Agent initialized with LLM: openai:gpt-4\")\nprint(\"Agent tools registered:\")\nprint(\"- search_vector_memory: Search for semantically similar content\")\nprint(\"- query_graph: Execute graph traversals from specific entities\")\nprint(\"- add_to_memory: Store new entities in the memory system\")\nprint(\"- link_entities: Create relationships between entities\")", - "id": "832d50545225edcf", - "outputs": [], - "execution_count": null + "id": "ex1", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:50.848619644Z", + "start_time": "2026-03-22T02:06:47.801159739Z" + } + }, + "source": "result = agent.run(\"What tools do you have available for managing my job search?\")\nprint(result.output)", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "I have the following tools available for managing your job search:\n", + "\n", + "1. **Search Memory**: I can search the memory for semantically similar content, allowing us to track existing applications, companies, contacts, and interactions.\n", + "\n", + "2. **Propose Memory Addition**: I can propose adding new content to memory, such as job postings, company information, or contact details, but it requires your confirmation before anything is stored.\n", + "\n", + "3. **Get Memory Context**: I can explore relationships between stored information via graph traversal, which can help us understand connections between companies, jobs, and contacts.\n", + "\n", + "These tools can help you keep track of your job applications and associated activities effectively. Let me know how you'd like to use them!\n" + ] + } + ], + "execution_count": 16 }, { "cell_type": "markdown", + "id": "ex2-header", "metadata": {}, "source": [ - "## Test the Agent" + "### Example 2: Adding Information to Memory\n", + "\n", + "The agent will propose adding items to memory. You need to confirm via `agent.confirm_memory_addition()`" + ] + }, + { + "cell_type": "code", + "id": "ex2-propose", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:57.567577238Z", + "start_time": "2026-03-22T02:06:52.570966533Z" + } + }, + "source": "# Use the trigger phrase to prompt memory storage\nresult = agent.run(\n \"I just applied to a Senior Software Engineer position at Google. \"\n \"The role is focused on distributed systems and will be based in Mountain View. \"\n \"Save this to memory.\"\n)\nprint(result.output)", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Yes, please add it to memory.\n" + ] + } + ], + "execution_count": 17 + }, + { + "cell_type": "code", + "id": "ex2-check", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:58.274527042Z", + "start_time": "2026-03-22T02:06:58.216435572Z" + } + }, + "source": [ + "# Check pending proposals\n", + "print(\"Pending proposals:\", agent.pending_proposals.keys())" ], - "id": "2bdc8896abdea4c3" + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Pending proposals: dict_keys(['42ec8dea-4ed8-4247-a5f1-d1abfb936d83'])\n" + ] + } + ], + "execution_count": 18 }, { "cell_type": "code", + "id": "ex2-confirm", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:06:59.520551614Z", + "start_time": "2026-03-22T02:06:58.907914374Z" + } + }, + "source": [ + "# Confirm the first proposal (if any exist)\n", + "if agent.pending_proposals:\n", + " proposal_id = list(agent.pending_proposals.keys())[0]\n", + " result = agent.confirm_memory_addition(proposal_id, action=\"add_new\")\n", + " print(result)\n", + "else:\n", + " print(\"No pending proposals\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Successfully added job to memory with ID: 9f57667f-21ae-43f6-baa6-bcb489602e3d\n" + ] + } + ], + "execution_count": 19 + }, + { + "cell_type": "markdown", + "id": "ex3-header", "metadata": {}, - "source": "# Initialize the memory system - choose your embedding model\n# Examples: 'openai:text-embedding-3-small', 'google-gla:text-embedding-004', 'cohere:embed-english-v3.0'\nembedding_model = infer_model('openai:text-embedding-3-small')\n\nmemory = VectorGraphMemory(\n qdrant_client=qdrant,\n janus_client=janus,\n embedding_model=embedding_model,\n collection_name=\"job_search_memory\"\n)\n\n# Create database context\ndb_context = DatabaseContext(memory=memory)\n\nprint(\"Memory system initialized with embedding model: openai:text-embedding-3-small\")\n\n# Run a simple query\nresult = await agent.run(\n \"What tools do you have available?\",\n deps=db_context\n)\n\nprint(result.data)", - "id": "7d8b3e4682ffbf71", - "outputs": [], - "execution_count": null + "source": [ + "### Example 3: Searching Memory" + ] + }, + { + "cell_type": "code", + "id": "ex3", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:04.808591525Z", + "start_time": "2026-03-22T02:07:01.485874320Z" + } + }, + "source": "# Test direct search with different thresholds\ntest_queries = [\n \"software engineering positions\",\n \"Google job application\",\n \"Senior Software Engineer\"\n]\n\nfor query in test_queries:\n print(f\"\\nQuery: '{query}'\")\n print(\"-\" * 50)\n \n # Try with lower threshold\n results = agent.store.search_similar_nodes(\n content=query,\n threshold=0.5, # Lower threshold to see what we get\n limit=5,\n project_id=\"job_search_2026\"\n )\n \n if results:\n for r in results:\n print(f\" Score: {r.similarity_score:.3f} | {r.content[:80]}...\")\n else:\n print(\" No results found\")\n\n# Now try with the agent\nprint(\"\\n\" + \"=\"*50)\nprint(\"Agent search:\")\nresult = agent.run(\"What software engineering positions have I applied to?\")\nprint(result.output)", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Query: 'software engineering positions'\n", + "--------------------------------------------------\n", + " Score: 0.521 | Applied to a Senior Software Engineer position at Google focused on distributed ...\n", + "\n", + "Query: 'Google job application'\n", + "--------------------------------------------------\n", + " Score: 0.524 | Applied to a Senior Software Engineer position at Google focused on distributed ...\n", + "\n", + "Query: 'Senior Software Engineer'\n", + "--------------------------------------------------\n", + " Score: 0.578 | Applied to a Senior Software Engineer position at Google focused on distributed ...\n", + "\n", + "==================================================\n", + "Agent search:\n", + "You have applied to a Senior Software Engineer position at Google, which focuses on distributed systems and is based in Mountain View. Would you like to add more details about this or check for more applications?\n" + ] + } + ], + "execution_count": 20 }, { "cell_type": "markdown", + "id": "ex4-header", "metadata": {}, "source": [ - "## Example: Job Search Use Case" + "### Example 4: Using Interval-Based Memory Trigger" + ] + }, + { + "cell_type": "code", + "id": "ex4-config", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:10.739169328Z", + "start_time": "2026-03-22T02:07:10.662956811Z" + } + }, + "source": [ + "# Change to interval mode - check every 5 messages\n", + "agent.set_memory_trigger(mode=\"interval\", message_interval=5)\n", + "print(\"✓ Memory trigger set to: every 5 messages\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Memory trigger set to: every 5 messages\n" + ] + } ], - "id": "9c0e7713f32af3b0" + "execution_count": 21 }, { "cell_type": "code", + "id": "ex4-test", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:28.576609464Z", + "start_time": "2026-03-22T02:07:11.929109499Z" + } + }, + "source": "# Send messages with meaningful content - on the 5th, agent should auto-check for memory items\nmessages = [\n \"I had a phone screening with Meta today\",\n \"The recruiter mentioned the team works on AI infrastructure\", \n \"They're looking for someone with distributed systems experience\",\n \"The position is remote-friendly with offices in Menlo Park\",\n \"I think this could be a good fit for my background\" # Message 5 - should trigger memory review\n]\n\nfor i, msg in enumerate(messages, 1):\n result = agent.run(msg)\n print(f\"Message {i}: {msg}\")\n print(f\"Response: {result.output}\\n\")\n\n# Check if any proposals were created\nif agent.pending_proposals:\n print(f\"\\n✓ Agent proposed {len(agent.pending_proposals)} memory addition(s)\")\n print(f\"Proposal IDs: {list(agent.pending_proposals.keys())}\")\nelse:\n print(\"\\n✗ No memory proposals generated\")", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Message 1: I had a phone screening with Meta today\n", + "Response: Should I add \"Phone screening with Meta\" to memory?\n", + "\n", + "Message 2: The recruiter mentioned the team works on AI infrastructure\n", + "Response: I found similar content related to your job application at Google. Here are the proposals for your confirmation:\n", + "\n", + "1. Should I add **\"The recruiter mentioned the team works on AI infrastructure.\"** as an interaction to memory?\n", + "2. Should I add **\"Google\"** as a company to memory?\n", + "\n", + "Please let me know how you'd like to proceed!\n", + "\n", + "Message 3: They're looking for someone with distributed systems experience\n", + "Response: I found a related entry for a Senior Software Engineer position at Google that focuses on distributed systems. Would you like to update this existing memory or create a new entry?\n", + "\n", + "Message 4: The position is remote-friendly with offices in Menlo Park\n", + "Response: I found a similar job application related to a Senior Software Engineer position at Google, which is based in Mountain View and likely to have a remote-friendly aspect. \n", + "\n", + "Would you like to provide more details about the new position for me to consider adding it to memory, or do you want to update the existing job application?\n", + "\n", + "Message 5: I think this could be a good fit for my background\n", + "Response: Could you please provide me with the details of the job posting or company you're referring to? That way, I can assist you in tracking it effectively.\n", + "\n", + "\n", + "✓ Agent proposed 3 memory addition(s)\n", + "Proposal IDs: ['0a2a5eda-73ca-4d1e-a437-d88bf327078f', '5920a404-c16b-428b-8b8d-2486cf57ab56', '1d9332a2-145a-428f-93d7-ca1a6aceb149']\n" + ] + } + ], + "execution_count": 22 + }, + { + "cell_type": "markdown", + "id": "ex5-header", "metadata": {}, "source": [ - "# Example: Add a job to memory\n", - "result = await agent.run(\n", - " \"Add a software engineer position at Google to my memory\",\n", - " deps=db_context\n", + "### Example 5: Dynamic Configuration Updates" + ] + }, + { + "cell_type": "code", + "id": "ex5", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:36.343561073Z", + "start_time": "2026-03-22T02:07:36.248256606Z" + } + }, + "source": [ + "# Update memory configuration on the fly\n", + "agent.configure_memory(\n", + " use_case_description=\"Track job search and networking contacts\",\n", + " memory_threshold_description=\"Store all interactions with recruiters and hiring managers\",\n", + " similarity_threshold=0.90, # Stricter duplicate detection\n", ")\n", - "\n", - "print(result.data)" + "print(\"✓ Memory configuration updated\")" ], - "id": "569ca3e38d410a94", - "outputs": [], - "execution_count": null + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Memory configuration updated\n" + ] + } + ], + "execution_count": 23 }, { - "cell_type": "code", + "cell_type": "markdown", + "id": "ex6-header", "metadata": {}, "source": [ - "# Example: Search for jobs\n", - "result = await agent.run(\n", - " \"Search for engineering positions in my memory\",\n", - " deps=db_context\n", - ")\n", + "### Example 6: Viewing Audit Logs" + ] + }, + { + "cell_type": "code", + "id": "ex6", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:41.863187715Z", + "start_time": "2026-03-22T02:07:41.750187519Z" + } + }, + "source": [ + "# Get recent audit entries\n", + "audit_entries = agent.get_audit_history(limit=10)\n", "\n", - "print(result.data)" + "print(f\"Recent audit log entries ({len(audit_entries)}):\")\n", + "for entry in audit_entries:\n", + " print(f\"\\n[{entry.timestamp}] {entry.operation_type}\")\n", + " print(f\" Summary: {entry.summary}\")\n", + " print(f\" Entities: {entry.affected_entities}\")" ], - "id": "7d404fff0c65b3f6", - "outputs": [], - "execution_count": null + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Recent audit log entries (5):\n", + "\n", + "[2026-03-22 01:55:04.470571] add_node\n", + " Summary: Added job: Applied to a Senior Software Engineer position at Google focused on distributed systems based in Mou\n", + " Entities: ['4f0ad4fa-9606-44ee-b0be-7b84e85297c0']\n", + "\n", + "[2026-03-22 01:56:42.333995] add_node\n", + " Summary: Added job: Applied to Senior Software Engineer position at Google focused on distributed systems, based in Moun\n", + " Entities: ['5a20cafc-5d1e-4e89-b7e6-5852f3d589ca']\n", + "\n", + "[2026-03-22 02:01:31.746646] add_node\n", + " Summary: Added job: Applied to a Senior Software Engineer position at Google focused on distributed systems, based in Mo\n", + " Entities: ['caf12bb9-2d54-4825-b000-f057385f8d8a']\n", + "\n", + "[2026-03-22 02:04:40.242716] add_node\n", + " Summary: Added job: Applied to a Senior Software Engineer position at Google focused on distributed systems based in Mou\n", + " Entities: ['a8b6948d-af92-4ef8-827a-4455076ad190']\n", + "\n", + "[2026-03-22 02:06:59.376155] add_node\n", + " Summary: Added job: Applied to a Senior Software Engineer position at Google focused on distributed systems, based in Mo\n", + " Entities: ['9f57667f-21ae-43f6-baa6-bcb489602e3d']\n" + ] + } + ], + "execution_count": 24 }, { "cell_type": "markdown", + "id": "cleanup-header", "metadata": {}, "source": [ "## Cleanup" - ], - "id": "3da5d52101dd7132" + ] }, { "cell_type": "code", - "metadata": {}, + "id": "cleanup", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:46.291797666Z", + "start_time": "2026-03-22T02:07:46.031868754Z" + } + }, "source": [ - "# Close connections\n", + "# Close JanusGraph connection\n", "janus.close()\n", - "print(\"Database connections closed\")" + "print(\"✓ Database connections closed\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Database connections closed\n" + ] + } ], - "id": "c28d83a09c52c4b4", + "execution_count": 25 + }, + { + "metadata": {}, + "cell_type": "code", "outputs": [], - "execution_count": null + "execution_count": null, + "source": "", + "id": "c972083aa8259ba1" } ], "metadata": { @@ -187,9 +552,9 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.14.0" + "version": "3.11.0" } }, "nbformat": 4, - "nbformat_minor": 4 -} \ No newline at end of file + "nbformat_minor": 5 +} diff --git a/playground_api.ipynb b/playground_api.ipynb new file mode 100644 index 0000000..4f64437 --- /dev/null +++ b/playground_api.ipynb @@ -0,0 +1,595 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "header", + "metadata": {}, + "source": [ + "# Vector Graph Memory API - Playground\n", + "\n", + "This notebook demonstrates using the Vector Graph Memory API for a job search use case.\n", + "\n", + "**Prerequisites:**\n", + "```bash\n", + "docker compose up -d\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "setup-header", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "id": "imports", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:45:24.046584906Z", + "start_time": "2026-03-22T02:45:23.952622740Z" + } + }, + "source": [ + "import requests\n", + "import json\n", + "\n", + "# API Configuration\n", + "API_BASE_URL = \"http://localhost:8000\"\n", + "SESSION_ID = \"playground-session-1\"\n", + "\n", + "def chat(message: str, session_id: str = SESSION_ID) -> dict:\n", + " \"\"\"Send a chat message to the API.\"\"\"\n", + " response = requests.post(\n", + " f\"{API_BASE_URL}/v1/chat/completions\",\n", + " json={\n", + " \"model\": \"vector-graph-memory\",\n", + " \"messages\": [{\"role\": \"user\", \"content\": message}],\n", + " \"user\": session_id,\n", + " }\n", + " )\n", + " response.raise_for_status()\n", + " return response.json()\n", + "\n", + "def get_proposals(session_id: str = SESSION_ID) -> dict:\n", + " \"\"\"Get pending memory proposals for the session.\"\"\"\n", + " response = requests.get(f\"{API_BASE_URL}/memory/proposals/{session_id}\")\n", + " response.raise_for_status()\n", + " return response.json()\n", + "\n", + "def confirm_proposal(proposal_id: str, session_id: str = SESSION_ID, action: str = \"add_new\") -> dict:\n", + " \"\"\"Confirm a memory proposal.\"\"\"\n", + " response = requests.post(\n", + " f\"{API_BASE_URL}/memory/confirm/{session_id}/{proposal_id}\",\n", + " params={\"action\": action}\n", + " )\n", + " response.raise_for_status()\n", + " return response.json()\n", + "\n", + "def get_audit_log(session_id: str = SESSION_ID) -> dict:\n", + " \"\"\"Get audit log for the session.\"\"\"\n", + " response = requests.get(f\"{API_BASE_URL}/memory/audit/{session_id}\")\n", + " response.raise_for_status()\n", + " return response.json()\n", + "\n", + "# Test API connection\n", + "response = requests.get(f\"{API_BASE_URL}/\")\n", + "if response.status_code == 200:\n", + " print(\"✓ Connected to Vector Graph Memory API\")\n", + " print(f\" Service: {response.json()['service']}\")\n", + " print(f\" Version: {response.json()['version']}\")\n", + "else:\n", + " print(\"✗ Failed to connect to API\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Connected to Vector Graph Memory API\n", + " Service: vector-graph-memory-api\n", + " Version: 1.0.0\n" + ] + } + ], + "execution_count": 1 + }, + { + "cell_type": "markdown", + "id": "ex1-header", + "metadata": {}, + "source": [ + "## Example 1: Basic Conversation" + ] + }, + { + "cell_type": "code", + "id": "ex1", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:45:30.036094562Z", + "start_time": "2026-03-22T02:45:26.473573190Z" + } + }, + "source": [ + "result = chat(\"What tools do you have available for managing my job search?\")\n", + "print(result[\"choices\"][0][\"message\"][\"content\"])" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "I have a few tools for managing your job search that involve memory management. Here are the key capabilities:\n", + "\n", + "1. **Search Memory**: I can search for existing relevant information from our previous conversations that might assist your job search.\n", + "2. **Propose Memory Addition**: I can suggest adding new information to our memory, such as job preferences or companies you're interested in, which can help keep track of your search progress.\n", + "3. **Get Memory Context**: I can retrieve context about previous conversations or memory items we've discussed to provide you a comprehensive view related to your job search.\n", + "\n", + "If there's anything specific you'd like me to assist with or retain in memory, just let me know!\n" + ] + } + ], + "execution_count": 2 + }, + { + "cell_type": "markdown", + "id": "ex2-header", + "metadata": {}, + "source": [ + "## Example 2: Adding Information to Memory\n", + "\n", + "The agent will propose adding items to memory when you use certain trigger phrases or patterns." + ] + }, + { + "cell_type": "code", + "id": "ex2-chat", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:45:39.675745887Z", + "start_time": "2026-03-22T02:45:33.963330050Z" + } + }, + "source": [ + "# Tell the agent about a job application\n", + "result = chat(\n", + " \"I just applied to a Senior Software Engineer position at Google. \"\n", + " \"The role is focused on distributed systems and will be based in Mountain View.\"\n", + ")\n", + "print(result[\"choices\"][0][\"message\"][\"content\"])" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Should I add \"Applied to a Senior Software Engineer position at Google focused on distributed systems, based in Mountain View\" to memory?\n" + ] + } + ], + "execution_count": 3 + }, + { + "cell_type": "code", + "id": "ex2-proposals", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:45:41.623045141Z", + "start_time": "2026-03-22T02:45:41.498823630Z" + } + }, + "source": [ + "# Check for pending proposals\n", + "proposals = get_proposals()\n", + "print(f\"\\nPending proposals: {len(proposals['proposals'])}\")\n", + "for proposal_id, proposal_data in proposals[\"proposals\"].items():\n", + " print(f\"\\nProposal ID: {proposal_id}\")\n", + " print(f\" Type: {proposal_data['entity_type']}\")\n", + " print(f\" Content: {proposal_data['content'][:100]}...\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Pending proposals: 1\n", + "\n", + "Proposal ID: 63710ea0-b5df-4737-93d6-2f0697855c10\n", + " Type: job\n", + " Content: Applied to a Senior Software Engineer position at Google focused on distributed systems, based in Mo...\n" + ] + } + ], + "execution_count": 4 + }, + { + "cell_type": "code", + "id": "ex2-confirm", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:45:44.671607280Z", + "start_time": "2026-03-22T02:45:43.183788111Z" + } + }, + "source": [ + "# Confirm the first proposal\n", + "proposals = get_proposals()\n", + "if proposals[\"proposals\"]:\n", + " proposal_id = list(proposals[\"proposals\"].keys())[0]\n", + " result = confirm_proposal(proposal_id)\n", + " print(result[\"message\"])\n", + "else:\n", + " print(\"No pending proposals\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Successfully added job to memory with ID: 2fc76ddc-b641-48c7-977b-07f733df290b\n" + ] + } + ], + "execution_count": 5 + }, + { + "cell_type": "markdown", + "id": "ex3-header", + "metadata": {}, + "source": "## Example 3: Querying Memory\n\nThe agent automatically searches memory when relevant to your queries.\n\n**Note:** For multi-turn conversations, you need to maintain conversation history across requests. Each API call includes the full conversation context so the agent understands references like \"the position\" or \"that job\"." + }, + { + "cell_type": "code", + "id": "ex3", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:46:02.433607278Z", + "start_time": "2026-03-22T02:45:50.965285400Z" + } + }, + "source": "# Ask about stored information - maintaining conversation history\nconversation = []\n\nqueries = [\n \"What software engineering positions have I applied to?\",\n \"Tell me about the Google job\",\n \"Where is the position located?\"\n]\n\nfor query in queries:\n print(f\"\\n{'='*60}\")\n print(f\"Q: {query}\")\n print(f\"{'='*60}\")\n\n # Add user message to conversation\n conversation.append({\"role\": \"user\", \"content\": query})\n\n # Send entire conversation history\n response = requests.post(\n f\"{API_BASE_URL}/v1/chat/completions\",\n json={\n \"model\": \"vector-graph-memory\",\n \"messages\": conversation,\n \"user\": SESSION_ID,\n }\n )\n result = response.json()\n\n # Add assistant response to conversation\n assistant_message = result[\"choices\"][0][\"message\"][\"content\"]\n conversation.append({\"role\": \"assistant\", \"content\": assistant_message})\n\n print(assistant_message)", + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "============================================================\n", + "Q: What software engineering positions have I applied to?\n", + "============================================================\n", + "It seems that you have applied to a Senior Software Engineer position at Google, which focuses on distributed systems and is based in Mountain View. \n", + "\n", + "Would you like to add any other specific details or updates related to your job applications, or should I add this position to your memory?\n", + "\n", + "============================================================\n", + "Q: Tell me about the Google job\n", + "============================================================\n", + "I found a similar entry in memory regarding your application for the Senior Software Engineer position at Google, which focuses on distributed systems and is based in Mountain View. \n", + "\n", + "Would you like to update the existing entry, or should I add this information as a new memory?\n", + "\n", + "============================================================\n", + "Q: Where is the position located?\n", + "============================================================\n", + "I found a similar entry in memory detailing your application for the Senior Software Engineer position at Google, which focuses on distributed systems and is based in Mountain View.\n", + "\n", + "Would you like to update this existing entry with any additional details or modifications, or should I keep it as is?\n" + ] + } + ], + "execution_count": 6 + }, + { + "cell_type": "markdown", + "id": "ex4-header", + "metadata": {}, + "source": [ + "## Example 4: Multi-Turn Conversation with Auto Memory\n", + "\n", + "The agent can automatically detect when to store information based on AI determination mode." + ] + }, + { + "cell_type": "code", + "id": "ex4", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:46:40.571087875Z", + "start_time": "2026-03-22T02:46:10.253855974Z" + } + }, + "source": [ + "# Conversation about a new opportunity\n", + "messages = [\n", + " \"I had a phone screening with Meta today\",\n", + " \"The recruiter mentioned the team works on AI infrastructure\",\n", + " \"They're looking for someone with distributed systems experience\",\n", + " \"The position is remote-friendly with offices in Menlo Park\",\n", + " \"I think this could be a good fit for my background\"\n", + "]\n", + "\n", + "for i, msg in enumerate(messages, 1):\n", + " print(f\"\\n--- Message {i} ---\")\n", + " print(f\"User: {msg}\")\n", + " result = chat(msg)\n", + " print(f\"Agent: {result['choices'][0]['message']['content']}\")\n", + "\n", + "# Check if any proposals were created\n", + "proposals = get_proposals()\n", + "print(f\"\\n{'='*60}\")\n", + "if proposals[\"proposals\"]:\n", + " print(f\"✓ Agent proposed {len(proposals['proposals'])} memory addition(s)\")\n", + " for proposal_id in proposals[\"proposals\"].keys():\n", + " print(f\" - {proposal_id}\")\n", + "else:\n", + " print(\"✗ No memory proposals generated\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "--- Message 1 ---\n", + "User: I had a phone screening with Meta today\n", + "Agent: Should I add \"User had a phone screening with Meta\" as an event to memory?\n", + "\n", + "--- Message 2 ---\n", + "User: The recruiter mentioned the team works on AI infrastructure\n", + "Agent: I have identified two additions to memory:\n", + "\n", + "1. Should I add **\"The team works on AI infrastructure\"** to memory as a team?\n", + "2. There is a similar item in memory regarding the **Senior Software Engineer position at Google**. Would you like me to add this as a new job, or update the existing one with the information about AI infrastructure?\n", + "\n", + "Please let me know how you’d like to proceed!\n", + "\n", + "--- Message 3 ---\n", + "User: They're looking for someone with distributed systems experience\n", + "Agent: I found a similar memory item related to applying for a Senior Software Engineer position at Google, which focused on distributed systems. Would you like to update this existing entry with the new information about them looking for someone with distributed systems experience, or should I create a new memory entry?\n", + "\n", + "--- Message 4 ---\n", + "User: The position is remote-friendly with offices in Menlo Park\n", + "Agent: Should I add \"The position is remote-friendly with offices in Menlo Park\" as a job to memory?\n", + "\n", + "--- Message 5 ---\n", + "User: I think this could be a good fit for my background\n", + "Agent: Should I add \"User believes a recent opportunity may be a good fit for their background\" to memory?\n", + "\n", + "============================================================\n", + "✓ Agent proposed 5 memory addition(s)\n", + " - d2944781-356d-43c6-a8ea-742c4c05f1d9\n", + " - 475ba242-0bd7-422f-b2ab-5c58a93254f1\n", + " - 56a57f91-8c45-4438-9da8-e175f6c198e3\n", + " - e3afa8e4-9e5d-4b42-8838-6f3923f299ce\n", + " - 727ef89f-2abb-4532-9fd5-3b4198a27fcd\n" + ] + } + ], + "execution_count": 7 + }, + { + "cell_type": "markdown", + "id": "ex5-header", + "metadata": {}, + "source": [ + "## Example 5: Viewing Audit History" + ] + }, + { + "cell_type": "code", + "id": "ex5", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:46:44.529193745Z", + "start_time": "2026-03-22T02:46:44.463448781Z" + } + }, + "source": [ + "# Get audit log for this session\n", + "audit = get_audit_log()\n", + "print(f\"Audit log entries for session '{SESSION_ID}': {len(audit['entries'])}\\n\")\n", + "\n", + "for entry in audit[\"entries\"]:\n", + " print(f\"[{entry['timestamp']}] {entry['operation']}\")\n", + " print(f\" {entry['summary']}\")\n", + " print(f\" Entities: {entry['entities']}\")\n", + " print()" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Audit log entries for session 'playground-session-1': 3\n", + "\n", + "[2026-03-22T02:33:30.768297] add_node\n", + " Added tools: The tools available for managing a job search include resources for job listings, resume building, i\n", + " Entities: ['51164d38-94fa-416c-9ee0-7ca0287d5b43']\n", + "\n", + "[2026-03-22T02:41:56.073909] add_node\n", + " Added job: Applied to a Senior Software Engineer position at Google focused on distributed systems based in Mou\n", + " Entities: ['e3451c47-61b8-4536-83ff-c71bbda5e3d9']\n", + "\n", + "[2026-03-22T02:45:44.449450] add_node\n", + " Added job: Applied to a Senior Software Engineer position at Google focused on distributed systems, based in Mo\n", + " Entities: ['2fc76ddc-b641-48c7-977b-07f733df290b']\n", + "\n" + ] + } + ], + "execution_count": 8 + }, + { + "cell_type": "markdown", + "id": "ex6-header", + "metadata": {}, + "source": [ + "## Example 6: Testing Different Query Types\n", + "\n", + "See how the agent handles various types of queries and determines what needs to be remembered." + ] + }, + { + "cell_type": "code", + "id": "ex6", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:47:14.180316140Z", + "start_time": "2026-03-22T02:46:51.142315888Z" + } + }, + "source": [ + "test_scenarios = [\n", + " {\n", + " \"name\": \"Factual Information\",\n", + " \"message\": \"The interview is scheduled for next Tuesday at 2pm\"\n", + " },\n", + " {\n", + " \"name\": \"Contact Information\",\n", + " \"message\": \"The recruiter's name is Sarah Johnson, email: sjohnson@meta.com\"\n", + " },\n", + " {\n", + " \"name\": \"Status Update\",\n", + " \"message\": \"I passed the technical interview!\"\n", + " },\n", + " {\n", + " \"name\": \"Casual Query\",\n", + " \"message\": \"What's the weather like?\"\n", + " },\n", + "]\n", + "\n", + "for scenario in test_scenarios:\n", + " print(f\"\\n{'='*60}\")\n", + " print(f\"Scenario: {scenario['name']}\")\n", + " print(f\"{'='*60}\")\n", + " print(f\"User: {scenario['message']}\")\n", + " \n", + " result = chat(scenario[\"message\"])\n", + " print(f\"Agent: {result['choices'][0]['message']['content']}\")\n", + " \n", + " # Check if a proposal was created\n", + " proposals_before = get_proposals()\n", + " print(f\"Proposals: {len(proposals_before['proposals'])}\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "============================================================\n", + "Scenario: Factual Information\n", + "============================================================\n", + "User: The interview is scheduled for next Tuesday at 2pm\n", + "Agent: Yes, please add \"Interview scheduled for next Tuesday at 2pm\" to memory.\n", + "Proposals: 6\n", + "\n", + "============================================================\n", + "Scenario: Contact Information\n", + "============================================================\n", + "User: The recruiter's name is Sarah Johnson, email: sjohnson@meta.com\n", + "Agent: I found two items that can be added to memory regarding Sarah Johnson: \n", + "\n", + "1. Recruiter: Sarah Johnson, Email: sjohnson@meta.com\n", + "2. Email of Sarah Johnson: sjohnson@meta.com (as a contact)\n", + "\n", + "Should I add both of these to memory?\n", + "Proposals: 8\n", + "\n", + "============================================================\n", + "Scenario: Status Update\n", + "============================================================\n", + "User: I passed the technical interview!\n", + "Agent: Yes, should I add \"Passed the technical interview\" as an event to memory?\n", + "Proposals: 9\n", + "\n", + "============================================================\n", + "Scenario: Casual Query\n", + "============================================================\n", + "User: What's the weather like?\n", + "Agent: It seems there are no relevant entries in memory related to the weather. If you have specific details or experiences about the weather you'd like me to remember, please let me know! Would you like to add something specific about the weather to memory?\n", + "Proposals: 9\n" + ] + } + ], + "execution_count": 9 + }, + { + "cell_type": "markdown", + "id": "helper-header", + "metadata": {}, + "source": [ + "## Helper Functions\n", + "\n", + "Additional utilities for working with the API." + ] + }, + { + "cell_type": "code", + "id": "helpers", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:47:25.726288497Z", + "start_time": "2026-03-22T02:47:25.656390570Z" + } + }, + "source": [ + "def print_response(result: dict):\n", + " \"\"\"Pretty print a chat response.\"\"\"\n", + " print(f\"Model: {result['model']}\")\n", + " print(f\"Message: {result['choices'][0]['message']['content']}\")\n", + " print(f\"Tokens: {result['usage']['total_tokens']} total\")\n", + " print(f\" (prompt: {result['usage']['prompt_tokens']}, completion: {result['usage']['completion_tokens']})\")\n", + "\n", + "def confirm_all_proposals(session_id: str = SESSION_ID):\n", + " \"\"\"Confirm all pending proposals for a session.\"\"\"\n", + " proposals = get_proposals(session_id)\n", + " confirmed = []\n", + " \n", + " for proposal_id in proposals[\"proposals\"].keys():\n", + " result = confirm_proposal(proposal_id, session_id)\n", + " confirmed.append((proposal_id, result[\"message\"]))\n", + " \n", + " return confirmed\n", + "\n", + "# Example usage\n", + "print(\"✓ Helper functions loaded\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Helper functions loaded\n" + ] + } + ], + "execution_count": 10 + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": "", + "id": "3a22c6e5a1516809" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 309ba7c28ebaaff1f7a7e05016dfebd644adbad6 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Tue, 24 Mar 2026 13:14:56 -0700 Subject: [PATCH 20/21] fix: resolve lint and type-check failures --- pyproject.toml | 16 +++++ scripts/init_janusgraph_schema.py | 11 ++-- src/vgm/MemoryAgent.py | 28 ++++++-- src/vgm/VectorGraphStore.py | 102 +++++++++++++++++++---------- src/vgm/api/server.py | 66 +++++++++++-------- src/vgm/audit/JSONLAuditBackend.py | 5 +- src/vgm/audit/MongoAuditBackend.py | 20 +++--- src/vgm/config/AuditConfig.py | 4 +- src/vgm/config/MemoryConfig.py | 8 ++- src/vgm/schemas/EdgeMetadata.py | 4 +- 10 files changed, 178 insertions(+), 86 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bffdc72..dea6da2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,9 +26,11 @@ api = [ dev = [ "ipykernel>=7.2.0", "jupyter>=1.1.1", + "mypy>=1.19.1", "notebook>=7.5.5", "pytest>=9.0.2", "pytest-cov>=7.1.0", + "ruff>=0.15.7", ] [build-system] @@ -37,3 +39,17 @@ build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] packages = ["src/vgm"] + +[tool.ruff] +extend-exclude = ["*.ipynb"] + +[tool.mypy] +python_version = "3.14" + +[[tool.mypy.overrides]] +module = [ + "gremlin_python", + "gremlin_python.*", + "nest_asyncio", +] +ignore_missing_imports = true diff --git a/scripts/init_janusgraph_schema.py b/scripts/init_janusgraph_schema.py index 34bca05..4ef7119 100755 --- a/scripts/init_janusgraph_schema.py +++ b/scripts/init_janusgraph_schema.py @@ -15,18 +15,18 @@ # Add src to path sys.path.insert(0, str(Path(__file__).parent.parent / "src")) -from gremlin_python.driver import client as gremlin_client +from gremlin_python.driver import client as gremlin_client # type: ignore[import-untyped] + def init_schema(): """Initialize JanusGraph schema.""" # Connect to JanusGraph - janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost') - janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182')) + janusgraph_host = os.getenv("JANUSGRAPH_HOST", "localhost") + janusgraph_port = int(os.getenv("JANUSGRAPH_PORT", "8182")) print(f"Connecting to JanusGraph at {janusgraph_host}:{janusgraph_port}...") janus = gremlin_client.Client( - f'ws://{janusgraph_host}:{janusgraph_port}/gremlin', - 'g' + f"ws://{janusgraph_host}:{janusgraph_port}/gremlin", "g" ) schema_script = """ @@ -81,5 +81,6 @@ def init_schema(): finally: janus.close() + if __name__ == "__main__": init_schema() diff --git a/src/vgm/MemoryAgent.py b/src/vgm/MemoryAgent.py index 6b08551..9feaaba 100644 --- a/src/vgm/MemoryAgent.py +++ b/src/vgm/MemoryAgent.py @@ -1,6 +1,5 @@ """Memory-enhanced AI agent with vector-graph storage.""" -import asyncio from typing import Optional, List, Dict, Any, Literal from uuid import uuid4 @@ -62,7 +61,9 @@ def __init__( trigger_config: Memory trigger configuration (optional) """ self.memory_config = memory_config - self.trigger_config = trigger_config or MemoryTriggerConfig(mode="phrase", trigger_phrase="check for memory items") + self.trigger_config = trigger_config or MemoryTriggerConfig( + mode="phrase", trigger_phrase="check for memory items" + ) self.audit_config = audit_config or AuditConfig() # Initialize storage layer @@ -220,7 +221,9 @@ def propose_memory_addition( proposal += f" {i}. [{sim.node_type}] {sim.content[:100]}... (similarity: {sim.similarity_score:.2f})\n" proposal += f"\nShould I add '{content[:100]}...' as a new {entity_type} to memory, or update one of the existing items?" else: - proposal = f"Should I add '{content[:100]}...' as a {entity_type} to memory?" + proposal = ( + f"Should I add '{content[:100]}...' as a {entity_type} to memory?" + ) return proposal @@ -252,7 +255,9 @@ def get_memory_context( formatted = [] for i, node in enumerate(results, 1): if isinstance(node, dict) and "content" in node: - formatted.append(f"{i}. [{node.get('node_type', 'unknown')}] {node['content'][:200]}...") + formatted.append( + f"{i}. [{node.get('node_type', 'unknown')}] {node['content'][:200]}..." + ) else: formatted.append(f"{i}. {node}") @@ -276,7 +281,9 @@ def configure_memory( if use_case_description: self.memory_config.use_case_description = use_case_description if memory_threshold_description: - self.memory_config.memory_threshold_description = memory_threshold_description + self.memory_config.memory_threshold_description = ( + memory_threshold_description + ) if similarity_threshold is not None: self.memory_config.similarity_threshold = similarity_threshold @@ -361,8 +368,14 @@ def confirm_memory_addition( project_id=self.memory_config.project_id, operation_type="add_node", summary=f"Added {proposal['entity_type']}: {proposal['content'][:100]}", - commands=[f"add_node(node_id={node_id})", *[f"add_edge(edge_id={eid})" for eid in edge_ids]], - metadata={"node_type": proposal["entity_type"], "relationship_count": len(edge_ids)}, + commands=[ + f"add_node(node_id={node_id})", + *[f"add_edge(edge_id={eid})" for eid in edge_ids], + ], + metadata={ + "node_type": proposal["entity_type"], + "relationship_count": len(edge_ids), + }, affected_entities=[node_id, *edge_ids], ) self.audit.log_operation(audit_entry) @@ -400,6 +413,7 @@ def confirm_memory_addition( except Exception as e: import traceback + error_details = traceback.format_exc() return f"Error executing memory operation: {str(e)}\n\nFull traceback:\n{error_details}" diff --git a/src/vgm/VectorGraphStore.py b/src/vgm/VectorGraphStore.py index 61ec4b5..c73de1a 100644 --- a/src/vgm/VectorGraphStore.py +++ b/src/vgm/VectorGraphStore.py @@ -1,18 +1,25 @@ """Vector and graph database storage layer.""" import asyncio -from typing import List, Optional, Dict, Any, Coroutine, TypeVar +from typing import Any, Callable, Coroutine, Dict, List, Optional, TypeVar, cast from datetime import datetime, timezone from qdrant_client import QdrantClient -from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue -from gremlin_python.driver import client as gremlin_client +from qdrant_client.models import ( + Distance, + VectorParams, + PointStruct, + Filter, + FieldCondition, + MatchValue, +) +from gremlin_python.driver import client as gremlin_client # type: ignore[import-untyped] from pydantic_ai import EmbeddingModel from .schemas import NodeMetadata, EdgeMetadata, SimilarNode from .config import VectorGraphConfig -T = TypeVar('T') +T = TypeVar("T") def _run_async(coro: Coroutine[Any, Any, T]) -> T: @@ -31,19 +38,20 @@ def _run_async(coro: Coroutine[Any, Any, T]) -> T: else: # There's a running loop # Check if it's patchable (uvloop isn't, but asyncio default loop is) - import nest_asyncio + import nest_asyncio # type: ignore[import-untyped] import threading + loop_type = type(loop).__name__ - if loop_type == 'Loop' and 'uvloop' in loop.__class__.__module__: + if loop_type == "Loop" and "uvloop" in loop.__class__.__module__: # uvloop can't be patched - run in new thread with new loop - result = None - exception = None + result: list[T] = [] + exception: Exception | None = None - def run_in_thread(): + def run_in_thread() -> None: nonlocal result, exception try: - result = asyncio.run(coro) + result.append(asyncio.run(coro)) except Exception as e: exception = e @@ -53,14 +61,16 @@ def run_in_thread(): if exception: raise exception - return result + if not result: + raise RuntimeError("Async operation did not produce a result") + return result[0] else: # Regular asyncio loop - can be patched for notebooks nest_asyncio.apply() return asyncio.run(coro) -def _run_in_thread(func, *args, **kwargs): +def _run_in_thread(func: Callable[..., T], *args: Any, **kwargs: Any) -> T: """Run a synchronous function in a separate thread. This is needed for Gremlin client operations which internally use @@ -68,13 +78,14 @@ def _run_in_thread(func, *args, **kwargs): (like FastAPI/uvicorn). """ import threading - result = None - exception = None - def run(): + result: list[T] = [] + exception: Exception | None = None + + def run() -> None: nonlocal result, exception try: - result = func(*args, **kwargs) + result.append(func(*args, **kwargs)) except Exception as e: exception = e @@ -84,7 +95,9 @@ def run(): if exception: raise exception - return result + if not result: + raise RuntimeError("Threaded operation did not produce a result") + return result[0] class VectorGraphStore: @@ -145,15 +158,16 @@ def _ensure_collection(self) -> None: collections = self.qdrant.get_collections().collections if not any(c.name == self.config.qdrant_collection for c in collections): # Get embedding dimension from a test embedding - test_result = _run_async(self.embedding_model.embed(["test"], input_type="document")) + test_result = _run_async( + self.embedding_model.embed(["test"], input_type="document") + ) embedding_dim = len(test_result.embeddings[0]) self.qdrant.create_collection( collection_name=self.config.qdrant_collection, vectors_config=VectorParams( - size=embedding_dim, - distance=Distance.COSINE - ) + size=embedding_dim, distance=Distance.COSINE + ), ) def _generate_embedding(self, text: str) -> List[float]: @@ -187,15 +201,18 @@ def search_similar_nodes( List of top K similar nodes, optionally filtered by threshold """ import logging + logger = logging.getLogger(__name__) embedding = self._generate_embedding(content) # Build search parameters - query_filter = None + query_filter: Filter | None = None if project_id: query_filter = Filter( - must=[FieldCondition(key="project_id", match=MatchValue(value=project_id))] + must=[ + FieldCondition(key="project_id", match=MatchValue(value=project_id)) + ] ) logger.info(f"[SEARCH] Query: '{content[:100]}...'") @@ -212,15 +229,23 @@ def search_similar_nodes( logger.info(f"[SEARCH] Found {len(response.points)} results") for i, result in enumerate(response.points, 1): - logger.info(f"[SEARCH] {i}. [{result.payload.get('node_type')}] score={result.score:.3f} content='{str(result.payload.get('content', ''))[:80]}...'") + payload = cast(Dict[str, Any], result.payload or {}) + logger.info( + f"[SEARCH] {i}. [{payload.get('node_type')}] " + f"score={result.score:.3f} content='{str(payload.get('content', ''))[:80]}...'" + ) return [ SimilarNode( node_id=str(result.id), - content=result.payload["content"], - node_type=result.payload["node_type"], + content=cast(str, (result.payload or {})["content"]), + node_type=cast(str, (result.payload or {})["node_type"]), similarity_score=result.score, - metadata={k: v for k, v in result.payload.items() if k not in ["content", "node_type"]}, + metadata={ + k: v + for k, v in cast(Dict[str, Any], result.payload or {}).items() + if k not in ["content", "node_type"] + }, ) for result in response.points ] @@ -241,6 +266,7 @@ def add_node( ValueError: If content exceeds max_content_size """ import logging + logger = logging.getLogger(__name__) # Check content size @@ -281,7 +307,7 @@ def add_node( ], ) - logger.info(f"[STORE] Successfully added to Qdrant") + logger.info("[STORE] Successfully added to Qdrant") # Store MINIMAL REFERENCE in JanusGraph (structure only) # Use generic "node" label, store actual type as a property @@ -324,7 +350,7 @@ def update_node( if not current: raise ValueError(f"Node {node_id} not found") - current_payload = current[0].payload + current_payload = dict(cast(Dict[str, Any], current[0].payload or {})) updated_at = datetime.now(timezone.utc).isoformat() # Update payload @@ -342,11 +368,16 @@ def update_node( embedding = self._generate_embedding(content) else: # Keep existing embedding - embedding = self.qdrant.retrieve( + stored_vector = self.qdrant.retrieve( collection_name=self.config.qdrant_collection, ids=[node_id], with_vectors=True, )[0].vector + if not isinstance(stored_vector, list): + raise TypeError(f"Expected stored vector list for node {node_id}") + if any(isinstance(value, list) for value in stored_vector): + raise TypeError(f"Expected flat vector list for node {node_id}") + embedding = [float(value) for value in cast(List[float], stored_vector)] if custom_metadata: current_payload.update(custom_metadata) @@ -424,7 +455,7 @@ def get_node(self, node_id: str) -> Optional[Dict[str, Any]]: if not results: return None - return results[0].payload + return cast(Dict[str, Any], results[0].payload or {}) def get_nodes_batch(self, node_ids: List[str]) -> Dict[str, Dict[str, Any]]: """Retrieve multiple nodes by ID from Qdrant. @@ -440,7 +471,10 @@ def get_nodes_batch(self, node_ids: List[str]) -> Dict[str, Dict[str, Any]]: ids=node_ids, ) - return {str(result.id): result.payload for result in results} + return { + str(result.id): cast(Dict[str, Any], result.payload or {}) + for result in results + } def traverse_from_node( self, @@ -463,7 +497,9 @@ def traverse_from_node( f"g.V().has('node_id', '{self._escape_gremlin_value(node_id)}')" f".{gremlin_steps}" ) - results = _run_in_thread(lambda: self.janus.submit(gremlin_query).all().result()) + results = _run_in_thread( + lambda: self.janus.submit(gremlin_query).all().result() + ) # If results contain node_ids, fetch full content from Qdrant # This is a simple heuristic - results might be node_ids or other data diff --git a/src/vgm/api/server.py b/src/vgm/api/server.py index f76dad5..bf83cbd 100644 --- a/src/vgm/api/server.py +++ b/src/vgm/api/server.py @@ -8,7 +8,7 @@ from fastapi import FastAPI, HTTPException from pydantic import BaseModel from qdrant_client import QdrantClient -from gremlin_python.driver import client as gremlin_client +from gremlin_python.driver import client as gremlin_client # type: ignore[import-untyped] from pydantic_ai.embeddings.openai import OpenAIEmbeddingModel from dotenv import load_dotenv @@ -22,14 +22,17 @@ # --- OpenAI-compatible API Models --- + class ChatMessage(BaseModel): """Chat message in OpenAI format.""" + role: str # system, user, assistant content: str class ChatCompletionRequest(BaseModel): """OpenAI chat completion request.""" + model: str messages: List[ChatMessage] temperature: Optional[float] = 0.7 @@ -40,6 +43,7 @@ class ChatCompletionRequest(BaseModel): class ChatCompletionChoice(BaseModel): """Single completion choice.""" + index: int message: ChatMessage finish_reason: str @@ -47,6 +51,7 @@ class ChatCompletionChoice(BaseModel): class ChatCompletionUsage(BaseModel): """Token usage statistics.""" + prompt_tokens: int completion_tokens: int total_tokens: int @@ -54,6 +59,7 @@ class ChatCompletionUsage(BaseModel): class ChatCompletionResponse(BaseModel): """OpenAI chat completion response.""" + id: str object: str = "chat.completion" created: int @@ -64,6 +70,7 @@ class ChatCompletionResponse(BaseModel): class ModelInfo(BaseModel): """Model information.""" + id: str object: str = "model" created: int @@ -72,14 +79,17 @@ class ModelInfo(BaseModel): class ModelList(BaseModel): """List of available models.""" + object: str = "list" data: List[ModelInfo] # --- Global state --- + class AppState: """Application state container.""" + agent: Optional[MemoryAgent] = None qdrant: Optional[QdrantClient] = None janus: Optional[gremlin_client.Client] = None @@ -91,6 +101,7 @@ class AppState: # --- Lifecycle management --- + @asynccontextmanager async def lifespan(_app: FastAPI): """Initialize and cleanup resources.""" @@ -100,73 +111,75 @@ async def lifespan(_app: FastAPI): # Configure logging logging.basicConfig( level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) print("🚀 Initializing Vector Graph Memory API...") # Connect to databases - qdrant_host = os.getenv('QDRANT_HOST', 'localhost') - qdrant_port = int(os.getenv('QDRANT_HTTP_PORT', '6333')) - janusgraph_host = os.getenv('JANUSGRAPH_HOST', 'localhost') - janusgraph_port = int(os.getenv('JANUSGRAPH_PORT', '8182')) + qdrant_host = os.getenv("QDRANT_HOST", "localhost") + qdrant_port = int(os.getenv("QDRANT_HTTP_PORT", "6333")) + janusgraph_host = os.getenv("JANUSGRAPH_HOST", "localhost") + janusgraph_port = int(os.getenv("JANUSGRAPH_PORT", "8182")) state.qdrant = QdrantClient(host=qdrant_host, port=qdrant_port) state.janus = gremlin_client.Client( - f'ws://{janusgraph_host}:{janusgraph_port}/gremlin', - 'g' + f"ws://{janusgraph_host}:{janusgraph_port}/gremlin", "g" ) print(f" ✓ Connected to Qdrant at {qdrant_host}:{qdrant_port}") print(f" ✓ Connected to JanusGraph at {janusgraph_host}:{janusgraph_port}") # Initialize memory agent - embedding_model = OpenAIEmbeddingModel('text-embedding-3-small') - llm_model = os.getenv('LLM_MODEL', 'openai:gpt-4o-mini') + embedding_model = OpenAIEmbeddingModel("text-embedding-3-small") + llm_model = os.getenv("LLM_MODEL", "openai:gpt-4o-mini") memory_config = MemoryConfig( use_case_description=os.getenv( - 'MEMORY_USE_CASE', - 'General purpose conversational memory' + "MEMORY_USE_CASE", "General purpose conversational memory" ), memory_threshold_description=os.getenv( - 'MEMORY_THRESHOLD', - 'Store important facts, decisions, and context from conversations' + "MEMORY_THRESHOLD", + "Store important facts, decisions, and context from conversations", ), - project_id=os.getenv('PROJECT_ID', 'default'), - similarity_threshold=float(os.getenv('SIMILARITY_THRESHOLD', '0.85')), + project_id=os.getenv("PROJECT_ID", "default"), + similarity_threshold=float(os.getenv("SIMILARITY_THRESHOLD", "0.85")), ) # Parse trigger mode with type safety - trigger_mode = os.getenv('TRIGGER_MODE', 'ai_determined') - valid_trigger_modes: tuple[Literal["phrase", "interval", "ai_determined"], ...] = ("phrase", "interval", "ai_determined") + trigger_mode = os.getenv("TRIGGER_MODE", "ai_determined") + valid_trigger_modes: tuple[Literal["phrase", "interval", "ai_determined"], ...] = ( + "phrase", + "interval", + "ai_determined", + ) if trigger_mode not in valid_trigger_modes: trigger_mode = "ai_determined" # Parse trigger interval with safety for empty strings - trigger_interval_str = os.getenv('TRIGGER_INTERVAL', '0') + trigger_interval_str = os.getenv("TRIGGER_INTERVAL", "0") trigger_interval = int(trigger_interval_str) if trigger_interval_str else 0 trigger_config = MemoryTriggerConfig( mode=cast(Literal["phrase", "interval", "ai_determined"], trigger_mode), - trigger_phrase=os.getenv('TRIGGER_PHRASE'), + trigger_phrase=os.getenv("TRIGGER_PHRASE"), message_interval=trigger_interval or None, ) # Parse audit backend with type safety - audit_backend = os.getenv('AUDIT_BACKEND', 'jsonl') + audit_backend = os.getenv("AUDIT_BACKEND", "jsonl") valid_audit_backends: tuple[Literal["jsonl", "mongodb"], ...] = ("jsonl", "mongodb") if audit_backend not in valid_audit_backends: audit_backend = "jsonl" audit_config = AuditConfig( backend=cast(Literal["jsonl", "mongodb"], audit_backend), - log_dir=os.getenv('AUDIT_LOG_DIR', '~/.vgm/logs'), + log_dir=os.getenv("AUDIT_LOG_DIR", "~/.vgm/logs"), ) system_prompt = os.getenv( - 'SYSTEM_PROMPT', - 'You are a helpful AI assistant with persistent memory capabilities.' + "SYSTEM_PROMPT", + "You are a helpful AI assistant with persistent memory capabilities.", ) state.agent = MemoryAgent( @@ -180,11 +193,11 @@ async def lifespan(_app: FastAPI): audit_config=audit_config, ) - print(f" ✓ Memory Agent initialized") + print(" ✓ Memory Agent initialized") print(f" - Model: {llm_model}") print(f" - Project: {memory_config.project_id}") print(f" - Trigger: {trigger_config.mode}") - print(f"✓ API ready on http://0.0.0.0:8000") + print("✓ API ready on http://0.0.0.0:8000") yield @@ -207,6 +220,7 @@ async def lifespan(_app: FastAPI): # --- Endpoints --- + @app.get("/") async def root(): """Health check endpoint.""" diff --git a/src/vgm/audit/JSONLAuditBackend.py b/src/vgm/audit/JSONLAuditBackend.py index bcc5427..c4c7108 100644 --- a/src/vgm/audit/JSONLAuditBackend.py +++ b/src/vgm/audit/JSONLAuditBackend.py @@ -48,7 +48,10 @@ def _rotate_log(self) -> None: # If rotated file exists, append to it; otherwise rename current if rotated_file.exists(): - with open(rotated_file, "a") as dest, open(self.current_log_file, "r") as src: + with ( + open(rotated_file, "a") as dest, + open(self.current_log_file, "r") as src, + ): dest.write(src.read()) self.current_log_file.unlink() else: diff --git a/src/vgm/audit/MongoAuditBackend.py b/src/vgm/audit/MongoAuditBackend.py index 24348be..111246a 100644 --- a/src/vgm/audit/MongoAuditBackend.py +++ b/src/vgm/audit/MongoAuditBackend.py @@ -1,6 +1,6 @@ """MongoDB-based audit backend.""" -from typing import List +from typing import Any, List from .AuditBackend import AuditBackend from ..schemas import AuditEntry @@ -12,8 +12,7 @@ class MongoAuditBackend(AuditBackend): def __init__(self, config: AuditConfig): try: - from pymongo import MongoClient, ASCENDING, IndexModel - from pymongo.errors import ConnectionFailure + from pymongo import MongoClient except ImportError: raise ImportError( "MongoDB backend requires pymongo. Install with: pip install vector-graph-memory[mongodb]" @@ -21,9 +20,11 @@ def __init__(self, config: AuditConfig): self.config = config if not config.connection_string: - raise ValueError("MongoDB backend requires connection_string in AuditConfig") + raise ValueError( + "MongoDB backend requires connection_string in AuditConfig" + ) - self.client = MongoClient(config.connection_string) + self.client: Any = MongoClient(config.connection_string) self.db = self.client[config.database] self.collection = self.db[config.collection] @@ -45,8 +46,7 @@ def _ensure_indexes(self) -> None: # Create TTL index if configured if self.config.ttl_days: self.collection.create_index( - "timestamp", - expireAfterSeconds=self.config.ttl_days * 24 * 60 * 60 + "timestamp", expireAfterSeconds=self.config.ttl_days * 24 * 60 * 60 ) def log_operation(self, entry: AuditEntry) -> None: @@ -66,9 +66,9 @@ def get_by_session(self, session_id: str) -> List[AuditEntry]: def get_entity_history(self, entity_id: str) -> List[AuditEntry]: """Get all operations affecting an entity.""" - docs = self.collection.find( - {"affected_entities": entity_id} - ).sort("timestamp", 1) + docs = self.collection.find({"affected_entities": entity_id}).sort( + "timestamp", 1 + ) return [AuditEntry.model_validate(doc) for doc in docs] def close(self) -> None: diff --git a/src/vgm/config/AuditConfig.py b/src/vgm/config/AuditConfig.py index c8f1fe9..7fb5934 100644 --- a/src/vgm/config/AuditConfig.py +++ b/src/vgm/config/AuditConfig.py @@ -15,7 +15,9 @@ class AuditConfig(BaseModel): rotation_period: Literal["daily", "weekly", "monthly"] = "monthly" # MongoDB-specific options - connection_string: Optional[str] = Field(default=None) # e.g., "mongodb://localhost:27017" + connection_string: Optional[str] = Field( + default=None + ) # e.g., "mongodb://localhost:27017" database: str = "vgm_audit" collection: str = "logs" ttl_days: Optional[int] = Field(default=None, gt=0) # Auto-expire old logs diff --git a/src/vgm/config/MemoryConfig.py b/src/vgm/config/MemoryConfig.py index 131e602..f04a7a0 100644 --- a/src/vgm/config/MemoryConfig.py +++ b/src/vgm/config/MemoryConfig.py @@ -7,6 +7,10 @@ class MemoryConfig(BaseModel): """Configuration for memory behavior and thresholds.""" use_case_description: str # Natural language: "Track job search activities..." - memory_threshold_description: str # Natural language: "Store specific job postings..." + memory_threshold_description: ( + str # Natural language: "Store specific job postings..." + ) project_id: str # Unique identifier for this memory project - similarity_threshold: float = Field(default=0.85, ge=0.0, le=1.0) # Duplicate detection threshold + similarity_threshold: float = Field( + default=0.85, ge=0.0, le=1.0 + ) # Duplicate detection threshold diff --git a/src/vgm/schemas/EdgeMetadata.py b/src/vgm/schemas/EdgeMetadata.py index c61ee4f..2dbead8 100644 --- a/src/vgm/schemas/EdgeMetadata.py +++ b/src/vgm/schemas/EdgeMetadata.py @@ -17,5 +17,7 @@ class EdgeMetadata(BaseModel): created_at: datetime = Field(default_factory=datetime.utcnow) source: str # session_id that created this edge project_id: str - confidence: Optional[float] = Field(default=None, ge=0.0, le=1.0) # Agent confidence in relationship + confidence: Optional[float] = Field( + default=None, ge=0.0, le=1.0 + ) # Agent confidence in relationship custom_metadata: Dict[str, Any] = Field(default_factory=dict) From 49765843cd2c5d2dc1b2d3d776b6a4af9cefd716 Mon Sep 17 00:00:00 2001 From: "Jason M. Cherry" Date: Tue, 24 Mar 2026 13:15:03 -0700 Subject: [PATCH 21/21] ci: install extras for pull request checks --- .github/workflows/ci.yml | 31 ++----------------------------- 1 file changed, 2 insertions(+), 29 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b8aa783..16e4fc8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,7 +22,7 @@ jobs: - name: Install dependencies run: | - uv sync + uv sync --all-extras - name: Run ruff check run: | @@ -47,35 +47,8 @@ jobs: - name: Install dependencies run: | - uv sync + uv sync --all-extras - name: Run mypy run: | uv run mypy . - - test: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.14' - - - name: Install uv - uses: astral-sh/setup-uv@v4 - - - name: Install dependencies - run: | - uv sync - - - name: Run pytest - run: | - uv run pytest --cov=vector_graph_memory --cov-report=xml --cov-report=term - - - name: Upload coverage to Codecov - uses: codecov/codecov-action@v4 - with: - file: ./coverage.xml - fail_ci_if_error: false