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/.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..16e4fc8 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,54 @@ +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 --all-extras + + - 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 --all-extras + + - name: Run mypy + run: | + uv run mypy . diff --git a/.gitignore b/.gitignore index b7faf40..b5958de 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. @@ -205,3 +205,7 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + + +## Custom Ignores +uv.lock 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. 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/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/README.md b/README.md new file mode 100644 index 0000000..43a259c --- /dev/null +++ b/README.md @@ -0,0 +1,176 @@ +# 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 + 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 + +### 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) + +## Development Setup + +### Prerequisites + +- Docker and Docker Compose +- Python 3.14+ +- uv (Python package manager) + +### Quick Start + +**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 +``` + +The API will be available at `http://localhost:8000` with OpenAI-compatible endpoints. + +**Option 2: Using the Python Library** + +```bash +# 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 +``` + +### 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 +jupyter notebook +``` + +## 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 + +The easiest way to test the API is with the included Open WebUI container: + +```bash +# Start all services including Open WebUI +docker compose up -d + +# Open WebUI will be available at http://localhost:3000 +``` + +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 + +Customize Open WebUI port in `.env`: + +```bash +WEBUI_PORT=3000 # Default port +``` + +See [API.md](API.md) for complete API documentation. + +## Status + +**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 + +See [LICENSE](LICENSE) file for details. diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..8a6abbc --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,119 @@ +services: + qdrant: + image: qdrant/qdrant:${QDRANT_VERSION:-latest} + container_name: ${QDRANT_CONTAINER_NAME:-vector-graph-qdrant} + expose: + - "6333" # HTTP API (internal only) + - "6334" # gRPC API (internal only) + 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:${JANUSGRAPH_VERSION:-latest} + container_name: ${JANUSGRAPH_CONTAINER_NAME:-vector-graph-janusgraph} + 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 + 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 + +volumes: + qdrant_storage: + driver: local + janusgraph_storage: + driver: local + mongodb_storage: + driver: local + open_webui_data: + driver: local + +networks: + vector-graph-network: + driver: bridge 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/playground.ipynb b/playground.ipynb new file mode 100644 index 0000000..059fc0f --- /dev/null +++ b/playground.ipynb @@ -0,0 +1,560 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "header", + "metadata": {}, + "source": [ + "# Vector Graph Memory - Playground\n", + "\n", + "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" + ] + }, + { + "cell_type": "code", + "id": "imports", + "metadata": { + "ExecuteTime": { + "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.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\n", + "\n", + "Make sure you've started the databases with `docker compose up -d`" + ] + }, + { + "cell_type": "code", + "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_HTTP_PORT', '6333'))\n", + "\n", + "qdrant = QdrantClient(host=qdrant_host, port=qdrant_port)\n", + "print(f\"✓ Connected to Qdrant at {qdrant_host}:{qdrant_port}\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Connected to Qdrant at localhost:6123\n" + ] + } + ], + "execution_count": 13 + }, + { + "cell_type": "code", + "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": [ + "## 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": [ + "## Usage Examples" + ] + }, + { + "cell_type": "markdown", + "id": "ex1-header", + "metadata": {}, + "source": [ + "### Example 1: Basic Conversation" + ] + }, + { + "cell_type": "code", + "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": [ + "### 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())" + ], + "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": [ + "### 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 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" + ] + } + ], + "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 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", + "print(\"✓ Memory configuration updated\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Memory configuration updated\n" + ] + } + ], + "execution_count": 23 + }, + { + "cell_type": "markdown", + "id": "ex6-header", + "metadata": {}, + "source": [ + "### 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(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}\")" + ], + "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" + ] + }, + { + "cell_type": "code", + "id": "cleanup", + "metadata": { + "ExecuteTime": { + "end_time": "2026-03-22T02:07:46.291797666Z", + "start_time": "2026-03-22T02:07:46.031868754Z" + } + }, + "source": [ + "# Close JanusGraph connection\n", + "janus.close()\n", + "print(\"✓ Database connections closed\")" + ], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Database connections closed\n" + ] + } + ], + "execution_count": 25 + }, + { + "metadata": {}, + "cell_type": "code", + "outputs": [], + "execution_count": null, + "source": "", + "id": "c972083aa8259ba1" + } + ], + "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.11.0" + } + }, + "nbformat": 4, + "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 +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..dea6da2 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,55 @@ +[project] +name = "vector-graph-memory" +version = "0.1.0" +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", + "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] +requires = ["hatchling"] +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 new file mode 100755 index 0000000..4ef7119 --- /dev/null +++ b/scripts/init_janusgraph_schema.py @@ -0,0 +1,86 @@ +#!/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 # 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")) + + 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() diff --git a/src/vgm/MemoryAgent.py b/src/vgm/MemoryAgent.py new file mode 100644 index 0000000..9feaaba --- /dev/null +++ b/src/vgm/MemoryAgent.py @@ -0,0 +1,510 @@ +"""Memory-enhanced AI agent with vector-graph storage.""" + +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..c73de1a --- /dev/null +++ b/src/vgm/VectorGraphStore.py @@ -0,0 +1,511 @@ +"""Vector and graph database storage layer.""" + +import asyncio +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 # type: ignore[import-untyped] +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 # type: ignore[import-untyped] + 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: list[T] = [] + exception: Exception | None = None + + def run_in_thread() -> None: + nonlocal result, exception + try: + result.append(asyncio.run(coro)) + except Exception as e: + exception = e + + thread = threading.Thread(target=run_in_thread) + thread.start() + thread.join() + + if exception: + raise exception + 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: 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 + run_until_complete() and fail when called from an existing event loop + (like FastAPI/uvicorn). + """ + import threading + + result: list[T] = [] + exception: Exception | None = None + + def run() -> None: + nonlocal result, exception + try: + result.append(func(*args, **kwargs)) + except Exception as e: + exception = e + + thread = threading.Thread(target=run) + thread.start() + thread.join() + + if exception: + raise exception + if not result: + raise RuntimeError("Threaded operation did not produce a result") + return result[0] + + +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: Filter | None = 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): + 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=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 cast(Dict[str, Any], result.payload or {}).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("[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 = dict(cast(Dict[str, Any], current[0].payload or {})) + 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 + 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) + 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 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. + + 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): cast(Dict[str, Any], result.payload or {}) + 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/__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/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..bf83cbd --- /dev/null +++ b/src/vgm/api/server.py @@ -0,0 +1,386 @@ +"""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 # type: ignore[import-untyped] +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(" ✓ Memory Agent initialized") + print(f" - Model: {llm_model}") + print(f" - Project: {memory_config.project_id}") + print(f" - Trigger: {trigger_config.mode}") + print("✓ 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/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..c4c7108 --- /dev/null +++ b/src/vgm/audit/JSONLAuditBackend.py @@ -0,0 +1,119 @@ +"""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..111246a --- /dev/null +++ b/src/vgm/audit/MongoAuditBackend.py @@ -0,0 +1,76 @@ +"""MongoDB-based audit backend.""" + +from typing import Any, 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 + 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: Any = 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", +] diff --git a/src/vgm/config/AuditConfig.py b/src/vgm/config/AuditConfig.py new file mode 100644 index 0000000..7fb5934 --- /dev/null +++ b/src/vgm/config/AuditConfig.py @@ -0,0 +1,23 @@ +"""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..f04a7a0 --- /dev/null +++ b/src/vgm/config/MemoryConfig.py @@ -0,0 +1,16 @@ +"""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", +] 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..2dbead8 --- /dev/null +++ b/src/vgm/schemas/EdgeMetadata.py @@ -0,0 +1,23 @@ +"""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", +] 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