first commit

.dockerignore

@@ -0,0 +1,16 @@
+.git
+.gitignore
+.env
+**/__pycache__
+**/*.pyc
+**/*.pyo
+**/*.pyd
+.Python
+env
+venv
+.venv
+chroma_db/
+ollama_data/
+.agent/
+.gemini/
+.specstory/

.env.example

@@ -0,0 +1,100 @@
+# API keys for different providers
+OPENAI_API_KEY=
+AZURE_OPENAI_API_KEY=
+DEEPSEEK_API_KEY=
+ANTHROPIC_API_KEY=
+GOOGLE_API_KEY=
+GROQ_API_KEY=
+OPENROUTER_API_KEY=
+USE_AWS_BEDROCK=false
+
+#Vertex AI
+GOOGLE_APPLICATION_CREDENTIALS=
+
+# Amazon Bedrock Knowledge Base ID
+AWS_KB_ID="<knowledge-base-id>"
+
+# Use a fake model for testing
+USE_FAKE_MODEL=false
+
+# Set a default model
+DEFAULT_MODEL=
+
+# If MODEL is set to "openai-compatible", set the following
+# This is just a flexible solution. If you need multiple model options, you still need to add it to models.py
+COMPATIBLE_MODEL=
+COMPATIBLE_API_KEY=
+COMPATIBLE_BASE_URL=
+
+# Web server configuration
+HOST=0.0.0.0
+PORT=7860
+
+# Authentication secret, HTTP bearer token header is required if set
+AUTH_SECRET=
+CORS_ORIGINS=http://localhost:3000,http://localhost:8081,http://localhost:5173
+
+# Langsmith configuration
+# LANGSMITH_TRACING=true
+# LANGSMITH_API_KEY=
+# LANGSMITH_PROJECT=default
+# LANGSMITH_ENDPOINT=https://api.smith.langchain.com
+
+# Application mode. If the value is "dev", it will enable uvicorn reload
+MODE=
+
+# Database type.
+# If the value is "postgres", then it will require Postgresql related environment variables.
+# If the value is "sqlite", then you can configure optional file path via SQLITE_DB_PATH
+DATABASE_TYPE=
+
+# If DATABASE_TYPE=sqlite (Optional)
+SQLITE_DB_PATH=
+
+# If DATABASE_TYPE=postgres
+# Docker Compose default values (will work with docker-compose setup)
+POSTGRES_USER=
+POSTGRES_PASSWORD=
+POSTGRES_HOST=
+POSTGRES_PORT=
+POSTGRES_DB=
+
+# you will be able to identify AST connections in Postgres Connection Manager under this Application Name
+# POSTGRES_APPLICATION_NAME = "agent-service-toolkit"
+# set these values to customize the number of connections in the pool. Saver and store have independent connection pools
+# POSTGRES_MIN_CONNECTIONS_PER_POOL=1
+# POSTGRES_MAX_CONNECTIONS_PER_POOL= 3
+
+# OpenWeatherMap API key
+OPENWEATHERMAP_API_KEY=
+
+# Add for running ollama
+# OLLAMA_MODEL=llama3.2
+# Note: set OLLAMA_BASE_URL if running service in docker and ollama on bare metal
+# OLLAMA_BASE_URL=http://host.docker.internal:11434
+
+# Add for running Azure OpenAI
+# AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com
+# AZURE_OPENAI_API_VERSION=2024-10-21
+# AZURE_OPENAI_DEPLOYMENT_MAP={"gpt-4o": "gpt-4o-deployment", "gpt-4o-mini": "gpt-4o-mini-deployment"}
+
+# Agent URL: used in Streamlit app - if not set, defaults to http://{HOST}:{PORT}
+# AGENT_URL=http://localhost:7860
+
+# LANGFUSE Configuration
+#LANGFUSE_TRACING=true
+#LANGFUSE_PUBLIC_KEY=pk-...
+#LANGFUSE_SECRET_KEY=sk-lf-....
+#LANGFUSE_HOST=http://localhost:3000
+
+# GitHub MCP Agent Configuration
+# GitHub Personal Access Token (required for GitHub MCP server)
+# If not set, the GitHub MCP agent will have no tools
+GITHUB_PAT=
+
+# Voice Features (Optional)
+# NOTE: Voice features are configured on the client (Streamlit app) side, not the server (API).
+# Requires OPENAI_API_KEY to be set (see above).
+# Set provider name to enable voice input/output. Leave empty to disable.
+VOICE_STT_PROVIDER=        # Speech-to-text provider (only 'openai' supported currently)
+VOICE_TTS_PROVIDER=        # Text-to-speech provider (only 'openai' supported currently)

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.github/workflows/sync-portfolio.yml

@@ -0,0 +1,70 @@
+name: Sync Notion Portfolio to PGVector Store
+
+on:
+  workflow_dispatch:
+    inputs:
+      since:
+        description: 'ISO 8601 date to sync from (e.g., 2024-01-01T00:00:00.000Z). Leave empty to use last sync date.'
+        required: false
+        type: string
+  schedule:
+    # Run daily at 2 AM UTC
+    - cron: "0 2 * * *"
+
+jobs:
+  sync-notion-portfolio-to-pgvector-store:
+    runs-on: ubuntu-latest
+    
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          version: "latest"
+
+      - name: Install dependencies
+        run: uv sync --frozen
+
+      - name: Sync Notion portfolio data to PGVector store
+        env:
+          # Notion API credentials
+          NOTION_TOKEN: ${{ secrets.NOTION_TOKEN }}
+          NOTION_EDUCATION_ID: ${{ secrets.NOTION_EDUCATION_ID }}
+          NOTION_EXPERIENCE_ID: ${{ secrets.NOTION_EXPERIENCE_ID }}
+          NOTION_PROJECT_ID: ${{ secrets.NOTION_PROJECT_ID }}
+          NOTION_TESTIMONIAL_ID: ${{ secrets.NOTION_TESTIMONIAL_ID }}
+          NOTION_BLOG_ID: ${{ secrets.NOTION_BLOG_ID }}
+          
+          # Database configuration
+          DATABASE_TYPE: postgres
+          POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
+          POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
+          POSTGRES_HOST: ${{ secrets.POSTGRES_HOST }}
+          POSTGRES_PORT: ${{ secrets.POSTGRES_PORT }}
+          POSTGRES_DB: ${{ secrets.POSTGRES_DB }}
+          
+          # LLM API keys (at least one required)
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+          GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+          GROQ_API_KEY: ${{ secrets.GROQ_API_KEY }}
+          
+          # Optional settings
+          OWNER: ${{ secrets.OWNER || 'Anuj Joshi' }}
+          VECTOR_STORE_COLLECTION_NAME: ${{ secrets.VECTOR_STORE_COLLECTION_NAME || 'portfolio' }}
+          
+          # Python path
+          PYTHONPATH: ${{ github.workspace }}/src
+        run: |
+          if [ -n "${{ inputs.since }}" ]; then
+            uv run python -m scripts.load_portfolio --since "${{ inputs.since }}"
+          else
+            uv run python -m scripts.load_portfolio
+          fi

.gitignore

@@ -0,0 +1,194 @@
+# Streamlit and sqlite
+.streamlit/secrets.toml
+checkpoints.db
+checkpoints.db-*
+
+# Langgraph
+.langgraph_api/
+
+# VSCode
+.vscode
+.DS_Store
+*.code-workspace
+
+# cursor
+.cursorindexingignore
+.specstory/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.python-version
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  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/
+
+# Data
+*.docx
+*.pdf
+
+# Chroma db
+chroma_db
+*.db
+*.sqlite3
+*.bin
+
+# Private Credentials, ignore everything in the folder but the .gitkeep file
+privatecredentials/*
+!privatecredentials/.gitkeep

Dockerfile

@@ -0,0 +1,53 @@
+# Stage 1: Get the official Ollama binary
+FROM ollama/ollama:latest AS ollama_source
+
+# Stage 2: Final Image
+FROM python:3.12.3-slim
+
+# Install system dependencies (curl for health, socat for port mapping)
+RUN apt-get update && apt-get install -y \
+    curl \
+    socat \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy Ollama from the official image
+COPY --from=ollama_source /usr/bin/ollama /usr/bin/ollama
+# Also copy the libraries which are necessary for Ollama to run
+COPY --from=ollama_source /usr/lib/ollama /usr/lib/ollama
+
+# Set up the non-root user
+RUN useradd -m -u 1000 user
+WORKDIR /app
+RUN chown user:user /app
+
+# Install uv and dependencies
+RUN pip install --no-cache-dir uv
+
+USER user
+COPY --chown=user:user pyproject.toml uv.lock ./
+RUN uv sync --frozen --no-dev --no-install-project
+
+# Copy app code
+COPY --chown=user:user src/agents/ ./agents/
+COPY --chown=user:user src/core/ ./core/
+COPY --chown=user:user src/memory/ ./memory/
+COPY --chown=user:user src/schema/ ./schema/
+COPY --chown=user:user src/service/ ./service/
+COPY --chown=user:user src/run_service.py ./run_service.py
+COPY --chown=user:user entrypoint.sh ./entrypoint.sh
+RUN chmod +x ./entrypoint.sh
+
+# Environment variables
+ENV HOST=0.0.0.0
+ENV PORT=7860
+ENV OLLAMA_HOST=0.0.0.0:11434
+ENV OLLAMA_KEEP_ALIVE=24h
+ENV OLLAMA_BASE_URL=http://localhost:11434
+
+EXPOSE 7860
+EXPOSE 11434
+
+RUN mkdir -p /home/user/.ollama
+VOLUME ["/home/user/.ollama"]
+
+CMD ["./entrypoint.sh"]

README.md

@@ -0,0 +1,240 @@
+---
+title: AI Agent Service Toolkit
+emoji: 🧰
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+pinned: false
+---
+
+# 🧰 AI Agent Service Toolkit
+
+[![build status](https://github.com/JoshuaC215/agent-service-toolkit/actions/workflows/test.yml/badge.svg)](https://github.com/JoshuaC215/agent-service-toolkit/actions/workflows/test.yml) [![codecov](https://codecov.io/github/JoshuaC215/agent-service-toolkit/graph/badge.svg?token=5MTJSYWD05)](https://codecov.io/github/JoshuaC215/agent-service-toolkit) [![Python Version](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2FJoshuaC215%2Fagent-service-toolkit%2Frefs%2Fheads%2Fmain%2Fpyproject.toml)](https://github.com/JoshuaC215/agent-service-toolkit/blob/main/pyproject.toml)
+[![GitHub License](https://img.shields.io/github/license/JoshuaC215/agent-service-toolkit)](https://github.com/JoshuaC215/agent-service-toolkit/blob/main/LICENSE) [![Streamlit App](https://static.streamlit.io/badges/streamlit_badge_black_red.svg)](https://agent-service-toolkit.streamlit.app/)
+
+A full toolkit for running an AI agent service built with LangGraph, FastAPI and Streamlit.
+
+It includes a [LangGraph](https://langchain-ai.github.io/langgraph/) agent, a [FastAPI](https://fastapi.tiangolo.com/) service to serve it, a client to interact with the service, and a [Streamlit](https://streamlit.io/) app that uses the client to provide a chat interface. Data structures and settings are built with [Pydantic](https://github.com/pydantic/pydantic).
+
+This project offers a template for you to easily build and run your own agents using the LangGraph framework. It demonstrates a complete setup from agent definition to user interface, making it easier to get started with LangGraph-based projects by providing a full, robust toolkit.
+
+**[🎥 Watch a video walkthrough of the repo and app](https://www.youtube.com/watch?v=pdYVHw_YCNY)**
+
+## Overview
+
+### [Try the app!](https://agent-service-toolkit.streamlit.app/)
+
+<a href="https://agent-service-toolkit.streamlit.app/"><img src="media/app_screenshot.png" width="600"></a>
+
+### Quickstart
+
+Run directly in python
+
+```sh
+# At least one LLM API key is required
+echo 'OPENAI_API_KEY=your_openai_api_key' >> .env
+
+# uv is the recommended way to install agent-service-toolkit, but "pip install ." also works
+# For uv installation options, see: https://docs.astral.sh/uv/getting-started/installation/
+curl -LsSf https://astral.sh/uv/0.7.19/install.sh | sh
+
+# Install dependencies. "uv sync" creates .venv automatically
+uv sync --frozen
+source .venv/bin/activate
+python src/run_service.py
+
+# In another shell
+source .venv/bin/activate
+streamlit run src/streamlit_app.py
+```
+
+Run with docker
+
+```sh
+echo 'OPENAI_API_KEY=your_openai_api_key' >> .env
+docker compose watch
+```
+
+### Architecture Diagram
+
+<img src="media/agent_architecture.png" width="600">
+
+### Key Features
+
+1. **LangGraph Agent and latest features**: A customizable agent built using the LangGraph framework. Implements the latest LangGraph v1.0 features including human in the loop with `interrupt()`, flow control with `Command`, long-term memory with `Store`, and `langgraph-supervisor`.
+1. **FastAPI Service**: Serves the agent with both streaming and non-streaming endpoints.
+1. **Advanced Streaming**: A novel approach to support both token-based and message-based streaming.
+1. **Streamlit Interface**: Provides a user-friendly chat interface for interacting with the agent, including voice input and output.
+1. **Multiple Agent Support**: Run multiple agents in the service and call by URL path. Available agents and models are described in `/info`
+1. **Asynchronous Design**: Utilizes async/await for efficient handling of concurrent requests.
+1. **Content Moderation**: Implements LlamaGuard for content moderation (requires Groq API key).
+1. **RAG Agent**: A basic RAG agent implementation using ChromaDB - see [docs](docs/RAG_Assistant.md).
+1. **Feedback Mechanism**: Includes a star-based feedback system integrated with LangSmith.
+1. **Docker Support**: Includes Dockerfiles and a docker compose file for easy development and deployment.
+1. **Testing**: Includes robust unit and integration tests for the full repo.
+
+### Key Files
+
+The repository is structured as follows:
+
+- `src/agents/`: Defines several agents with different capabilities
+- `src/schema/`: Defines the protocol schema
+- `src/core/`: Core modules including LLM definition and settings
+- `src/service/service.py`: FastAPI service to serve the agents
+- `src/client/client.py`: Client to interact with the agent service
+- `src/streamlit_app.py`: Streamlit app providing a chat interface
+- `tests/`: Unit and integration tests
+
+## Setup and Usage
+
+1. Clone the repository:
+
+   ```sh
+   git clone https://github.com/JoshuaC215/agent-service-toolkit.git
+   cd agent-service-toolkit
+   ```
+
+2. Set up environment variables:
+   Create a `.env` file in the root directory. At least one LLM API key or configuration is required. See the [`.env.example` file](./.env.example) for a full list of available environment variables, including a variety of model provider API keys, header-based authentication, LangSmith tracing, testing and development modes, and OpenWeatherMap API key.
+
+3. You can now run the agent service and the Streamlit app locally, either with Docker or just using Python. The Docker setup is recommended for simpler environment setup and immediate reloading of the services when you make changes to your code.
+
+### Additional setup for specific AI providers
+
+- [Setting up Ollama](docs/Ollama.md)
+- [Setting up VertexAI](docs/VertexAI.md)
+- [Setting up RAG with ChromaDB](docs/RAG_Assistant.md)
+
+### Building or customizing your own agent
+
+To customize the agent for your own use case:
+
+1. Add your new agent to the `src/agents` directory. You can copy `research_assistant.py` or `chatbot.py` and modify it to change the agent's behavior and tools.
+1. Import and add your new agent to the `agents` dictionary in `src/agents/agents.py`. Your agent can be called by `/<your_agent_name>/invoke` or `/<your_agent_name>/stream`.
+1. Adjust the Streamlit interface in `src/streamlit_app.py` to match your agent's capabilities.
+
+
+### Handling Private Credential files
+
+If your agents or chosen LLM require file-based credential files or certificates, the `privatecredentials/` has been provided for your development convenience. All contents, excluding the `.gitkeep` files, are ignored by git and docker's build process. See [Working with File-based Credentials](docs/File_Based_Credentials.md) for suggested use.
+
+
+### Docker Setup
+
+This project includes a Docker setup for easy development and deployment. The `compose.yaml` file defines three services: `postgres`, `agent_service` and `streamlit_app`. The `Dockerfile` for each service is in their respective directories.
+
+For local development, we recommend using [docker compose watch](https://docs.docker.com/compose/file-watch/). This feature allows for a smoother development experience by automatically updating your containers when changes are detected in your source code.
+
+1. Make sure you have Docker and Docker Compose (>= [v2.23.0](https://docs.docker.com/compose/release-notes/#2230)) installed on your system.
+
+2. Create a `.env` file from the `.env.example`. At minimum, you need to provide an LLM API key (e.g., OPENAI_API_KEY).
+   ```sh
+   cp .env.example .env
+   # Edit .env to add your API keys
+   ```
+
+3. Build and launch the services in watch mode:
+
+   ```sh
+   docker compose watch
+   ```
+
+   This will automatically:
+   - Start a PostgreSQL database service that the agent service connects to
+   - Start the agent service with FastAPI
+   - Start the Streamlit app for the user interface
+
+4. The services will now automatically update when you make changes to your code:
+   - Changes in the relevant python files and directories will trigger updates for the relevant services.
+   - NOTE: If you make changes to the `pyproject.toml` or `uv.lock` files, you will need to rebuild the services by running `docker compose up --build`.
+
+5. Access the Streamlit app by navigating to `http://localhost:8501` in your web browser.
+
+6. The agent service API will be available at `http://0.0.0.0:7860`. You can also use the OpenAPI docs at `http://0.0.0.0:7860/redoc`.
+
+7. Use `docker compose down` to stop the services.
+
+This setup allows you to develop and test your changes in real-time without manually restarting the services.
+
+### Building other apps on the AgentClient
+
+The repo includes a generic `src/client/client.AgentClient` that can be used to interact with the agent service. This client is designed to be flexible and can be used to build other apps on top of the agent. It supports both synchronous and asynchronous invocations, and streaming and non-streaming requests.
+
+See the `src/run_client.py` file for full examples of how to use the `AgentClient`. A quick example:
+
+```python
+from client import AgentClient
+client = AgentClient()
+
+response = client.invoke("Tell me a brief joke?")
+response.pretty_print()
+# ================================== Ai Message ==================================
+#
+# A man walked into a library and asked the librarian, "Do you have any books on Pavlov's dogs and Schrödinger's cat?"
+# The librarian replied, "It rings a bell, but I'm not sure if it's here or not."
+
+```
+
+### Development with LangGraph Studio
+
+The agent supports [LangGraph Studio](https://langchain-ai.github.io/langgraph/concepts/langgraph_studio/), the IDE for developing agents in LangGraph.
+
+`langgraph-cli[inmem]` is installed with `uv sync`. You can simply add your `.env` file to the root directory as described above, and then launch LangGraph Studio with `langgraph dev`. Customize `langgraph.json` as needed. See the [local quickstart](https://langchain-ai.github.io/langgraph/cloud/how-tos/studio/quick_start/#local-development-server) to learn more.
+
+### Local development without Docker
+
+You can also run the agent service and the Streamlit app locally without Docker, just using a Python virtual environment.
+
+1. Create a virtual environment and install dependencies:
+
+   ```sh
+   uv sync --frozen
+   source .venv/bin/activate
+   ```
+
+2. Run the FastAPI server:
+
+   ```sh
+   python src/run_service.py
+   ```
+
+3. In a separate terminal, run the Streamlit app:
+
+   ```sh
+   streamlit run src/streamlit_app.py
+   ```
+
+4. Open your browser and navigate to the URL provided by Streamlit (usually `http://localhost:8501`).
+
+## Projects built with or inspired by agent-service-toolkit
+
+The following are a few of the public projects that drew code or inspiration from this repo.
+
+- **[PolyRAG](https://github.com/QuentinFuxa/PolyRAG)** - Extends agent-service-toolkit with RAG capabilities over both PostgreSQL databases and PDF documents.
+- **[alexrisch/agent-web-kit](https://github.com/alexrisch/agent-web-kit)** - A Next.JS frontend for agent-service-toolkit
+- **[raushan-in/dapa](https://github.com/raushan-in/dapa)** - Digital Arrest Protection App (DAPA) enables users to report financial scams and frauds efficiently via a user-friendly platform.
+
+**Please create a pull request editing the README or open a discussion with any new ones to be added!** Would love to include more projects.
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request. Currently the tests need to be run using the local development without Docker setup. To run the tests for the agent service:
+
+1. Ensure you're in the project root directory and have activated your virtual environment.
+
+2. Install the development dependencies and pre-commit hooks:
+
+   ```sh
+   uv sync --frozen
+   pre-commit install
+   ```
+
+3. Run the tests using pytest:
+
+   ```sh
+   pytest
+   ```
+
+## License
+
+This project is licensed under the MIT License - see the LICENSE file for details.

compose.yaml

@@ -0,0 +1,35 @@
+services:
+  portfolio_bot:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    ports:
+      - "7860:7860"
+    env_file:
+      - .env
+    environment:
+      - OLLAMA_BASE_URL=http://localhost:11434
+      - OLLAMA_MODEL=qwen3:0.6b
+      - OLLAMA_EMBEDDING_MODEL=embeddinggemma:300m
+    healthcheck:
+      test: [ "CMD", "curl", "-f", "http://localhost:7860/health" ]
+      interval: 30s
+      timeout: 10s
+      retries: 5
+    develop:
+      watch:
+        - path: src/agents/
+          action: sync+restart
+          target: /app/agents/
+        - path: src/schema/
+          action: sync+restart
+          target: /app/schema/
+        - path: src/service/
+          action: sync+restart
+          target: /app/service/
+        - path: src/core/
+          action: sync+restart
+          target: /app/core/
+        - path: src/memory/
+          action: sync+restart
+          target: /app/memory/

entrypoint.sh

@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+set -e
+
+echo "Starting Ollama..."
+ollama serve &
+OLLAMA_PID=$!
+
+# Wait for Ollama to start
+until curl -s http://localhost:11434/api/tags > /dev/null; do
+    echo "Waiting for Ollama..."
+    sleep 2
+done
+
+echo "Pulling models..."
+# Pull models based on environment variables or defaults
+OLLAMA_MODEL=${OLLAMA_MODEL:-"qwen3:0.6b"}
+OLLAMA_EMBEDDING_MODEL=${OLLAMA_EMBEDDING_MODEL:-"embeddinggemma:300m"}
+
+echo "Pulling model: $OLLAMA_MODEL"
+ollama pull $OLLAMA_MODEL
+
+echo "Pulling embedding model: $OLLAMA_EMBEDDING_MODEL"
+ollama pull $OLLAMA_EMBEDDING_MODEL
+
+echo "Ollama ready."
+
+echo "Starting Backend Service on port 7860..."
+# Set variables for the python service
+export HOST=0.0.0.0
+export PORT=7860
+export OLLAMA_BASE_URL=http://localhost:11434
+
+# Ensure logs are visible
+export PYTHONUNBUFFERED=1
+
+# Start the python service
+uv run python run_service.py

pyproject.toml

@@ -0,0 +1,108 @@
+[project]
+name = "agent-service-toolkit"
+version = "0.1.0"
+description = "Full toolkit for running an AI agent service built with LangGraph, FastAPI and Streamlit"
+readme = "README.md"
+authors = [{ name = "Joshua Carroll", email = "carroll.joshk@gmail.com" }]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "License :: OSI Approved :: MIT License",
+    "Framework :: FastAPI",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+
+requires-python = ">=3.11,<3.14"
+
+dependencies = [
+    "docx2txt ~=0.8",
+    "duckduckgo-search>=7.3.0",
+    "fastapi ~=0.115.5",
+    "grpcio >=1.68.0",
+    "httpx ~=0.28.0",
+    "jiter ~=0.8.2",
+    "langchain ~=1.0.5",
+    "langchain-core ~=1.0.0",
+    "langchain-community ~=0.4.1",
+    "langchain-anthropic ~=1.0.0",
+    "langchain-aws ~=1.0.0",
+    "langchain-postgres ~=0.0.9",
+    "langchain-google-genai ~=3.0.0",
+    "langchain-google-vertexai>=3.0.3",
+    "langchain-groq ~=1.0.1",
+    "langchain-ollama ~=1.0.0",
+    "langchain-openai ~=1.0.2",
+    "langfuse >=2.65.0",
+    "langgraph ~=1.0.0",
+    "langgraph-checkpoint-mongodb ~=0.1.3",
+    "langgraph-checkpoint-postgres ~=2.0.13",
+    "langgraph-checkpoint-sqlite ~=2.0.1",
+    "langgraph-supervisor ~=0.0.31",
+    "langsmith ~=0.4.0",
+    "numexpr ~=2.10.1",
+    "numpy ~=2.3.4",
+    "onnxruntime ~= 1.21.1",
+    "pandas ~=2.2.3",
+    "psycopg[binary,pool] ~=3.2.4",
+    "pyarrow >=18.1.0",
+    "pydantic ~=2.10.1",
+    "pydantic-settings ~=2.12.0",
+    "pypdf ~=5.3.0",
+    "pyowm ~=3.3.0",
+    "python-dotenv ~=1.0.1",
+    "setuptools ~=75.6.0",
+    "streamlit ~=1.52.0",
+    "tiktoken >=0.8.0",
+    "uvicorn ~=0.32.1",
+    "langchain-mcp-adapters>=0.1.10",
+    "ddgs>=9.9.1",
+]
+
+[dependency-groups]
+dev = [
+    "langgraph-cli[inmem]",
+    "pre-commit",
+    "pytest",
+    "pytest-cov",
+    "pytest-env",
+    "pytest-asyncio",
+    "ruff",
+    "mypy",
+]
+
+# Group for the minimal dependencies to run just the client and Streamlit app.
+# These are also installed in the default dependencies.
+# To install run: `uv sync --frozen --only-group client`
+client = [
+    "httpx~=0.28.0",
+    "pydantic ~=2.10.1",
+    "python-dotenv ~=1.0.1",
+    "streamlit~=1.52.0",
+]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+extend-select = ["I", "U"]
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+asyncio_default_fixture_loop_scope = "function"
+
+[tool.pytest_env]
+OPENAI_API_KEY = "sk-fake-openai-key"
+
+[tool.mypy]
+plugins = "pydantic.mypy"
+exclude = "src/streamlit_app.py"
+
+[[tool.mypy.overrides]]
+module = ["numexpr.*"]
+follow_untyped_imports = true
+
+[[tool.mypy.overrides]]
+module = ["langchain_mcp_adapters.*"]
+ignore_missing_imports = true

src/agents/__init__.py

@@ -0,0 +1,17 @@
+from agents.agents import (
+    DEFAULT_AGENT,
+    AgentGraph,
+    AgentGraphLike,
+    get_agent,
+    get_all_agent_info,
+    load_agent,
+)
+
+__all__ = [
+    "get_agent",
+    "load_agent",
+    "get_all_agent_info",
+    "DEFAULT_AGENT",
+    "AgentGraph",
+    "AgentGraphLike",
+]

src/agents/agents.py

@@ -0,0 +1,97 @@
+from dataclasses import dataclass
+
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.pregel import Pregel
+
+from agents.bg_task_agent.bg_task_agent import bg_task_agent
+from agents.chatbot import chatbot
+from agents.command_agent import command_agent
+from agents.github_mcp_agent.github_mcp_agent import github_mcp_agent
+from agents.interrupt_agent import interrupt_agent
+from agents.knowledge_base_agent import kb_agent
+from agents.langgraph_supervisor_agent import langgraph_supervisor_agent
+from agents.langgraph_supervisor_hierarchy_agent import langgraph_supervisor_hierarchy_agent
+from agents.portfolio_agent.portfolio_agent import portfolio_agent
+
+from agents.lazy_agent import LazyLoadingAgent
+from agents.rag_assistant import rag_assistant
+from agents.research_assistant import research_assistant
+from schema import AgentInfo
+
+DEFAULT_AGENT = "portfolio-agent"
+
+# Type alias to handle LangGraph's different agent patterns
+# - @entrypoint functions return Pregel
+# - StateGraph().compile() returns CompiledStateGraph
+AgentGraph = CompiledStateGraph | Pregel  # What get_agent() returns (always loaded)
+AgentGraphLike = CompiledStateGraph | Pregel | LazyLoadingAgent  # What can be stored in registry
+
+
+@dataclass
+class Agent:
+    description: str
+    graph_like: AgentGraphLike
+
+
+agents: dict[str, Agent] = {
+    "chatbot": Agent(description="A simple chatbot.", graph_like=chatbot),
+    "research-assistant": Agent(
+        description="A research assistant with web search and calculator.",
+        graph_like=research_assistant,
+    ),
+    "rag-assistant": Agent(
+        description="A RAG assistant with access to information in a database.",
+        graph_like=rag_assistant,
+    ),
+    "portfolio-agent": Agent(
+        description="A portfolio assistant with access to information in a database.",
+        graph_like=portfolio_agent,
+    ),
+    "command-agent": Agent(description="A command agent.", graph_like=command_agent),
+    "bg-task-agent": Agent(description="A background task agent.", graph_like=bg_task_agent),
+    "langgraph-supervisor-agent": Agent(
+        description="A langgraph supervisor agent", graph_like=langgraph_supervisor_agent
+    ),
+    "langgraph-supervisor-hierarchy-agent": Agent(
+        description="A langgraph supervisor agent with a nested hierarchy of agents",
+        graph_like=langgraph_supervisor_hierarchy_agent,
+    ),
+    "interrupt-agent": Agent(
+        description="An agent the uses interrupts.", graph_like=interrupt_agent
+    ),
+    "knowledge-base-agent": Agent(
+        description="A retrieval-augmented generation agent using Amazon Bedrock Knowledge Base",
+        graph_like=kb_agent,
+    ),
+    "github-mcp-agent": Agent(
+        description="A GitHub agent with MCP tools for repository management and development workflows.",
+        graph_like=github_mcp_agent,
+    ),
+}
+
+
+async def load_agent(agent_id: str) -> None:
+    """Load lazy agents if needed."""
+    graph_like = agents[agent_id].graph_like
+    if isinstance(graph_like, LazyLoadingAgent):
+        await graph_like.load()
+
+
+def get_agent(agent_id: str) -> AgentGraph:
+    """Get an agent graph, loading lazy agents if needed."""
+    agent_graph = agents[agent_id].graph_like
+
+    # If it's a lazy loading agent, ensure it's loaded and return its graph
+    if isinstance(agent_graph, LazyLoadingAgent):
+        if not agent_graph._loaded:
+            raise RuntimeError(f"Agent {agent_id} not loaded. Call load() first.")
+        return agent_graph.get_graph()
+
+    # Otherwise return the graph directly
+    return agent_graph
+
+
+def get_all_agent_info() -> list[AgentInfo]:
+    return [
+        AgentInfo(key=agent_id, description=agent.description) for agent_id, agent in agents.items()
+    ]

src/agents/bg_task_agent/bg_task_agent.py

@@ -0,0 +1,62 @@
+import asyncio
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnableSerializable
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.types import StreamWriter
+
+from agents.bg_task_agent.task import Task
+from core import get_model, settings
+
+
+class AgentState(MessagesState, total=False):
+    """`total=False` is PEP589 specs.
+
+    documentation: https://typing.readthedocs.io/en/latest/spec/typeddict.html#totality
+    """
+
+
+def wrap_model(model: BaseChatModel) -> RunnableSerializable[AgentState, AIMessage]:
+    preprocessor = RunnableLambda(
+        lambda state: state["messages"],
+        name="StateModifier",
+    )
+    return preprocessor | model  # type: ignore[return-value]
+
+
+async def acall_model(state: AgentState, config: RunnableConfig) -> AgentState:
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(m)
+    response = await model_runnable.ainvoke(state, config)
+
+    # We return a list, because this will get added to the existing list
+    return {"messages": [response]}
+
+
+async def bg_task(state: AgentState, writer: StreamWriter) -> AgentState:
+    task1 = Task("Simple task 1...", writer)
+    task2 = Task("Simple task 2...", writer)
+
+    task1.start()
+    await asyncio.sleep(2)
+    task2.start()
+    await asyncio.sleep(2)
+    task1.write_data(data={"status": "Still running..."})
+    await asyncio.sleep(2)
+    task2.finish(result="error", data={"output": 42})
+    await asyncio.sleep(2)
+    task1.finish(result="success", data={"output": 42})
+    return {"messages": []}
+
+
+# Define the graph
+agent = StateGraph(AgentState)
+agent.add_node("model", acall_model)
+agent.add_node("bg_task", bg_task)
+agent.set_entry_point("bg_task")
+
+agent.add_edge("bg_task", "model")
+agent.add_edge("model", END)
+
+bg_task_agent = agent.compile()

src/agents/bg_task_agent/task.py

@@ -0,0 +1,53 @@
+from typing import Literal
+from uuid import uuid4
+
+from langchain_core.messages import BaseMessage
+from langgraph.types import StreamWriter
+
+from agents.utils import CustomData
+from schema.task_data import TaskData
+
+
+class Task:
+    def __init__(self, task_name: str, writer: StreamWriter | None = None) -> None:
+        self.name = task_name
+        self.id = str(uuid4())
+        self.state: Literal["new", "running", "complete"] = "new"
+        self.result: Literal["success", "error"] | None = None
+        self.writer = writer
+
+    def _generate_and_dispatch_message(self, writer: StreamWriter | None, data: dict):
+        writer = writer or self.writer
+        task_data = TaskData(name=self.name, run_id=self.id, state=self.state, data=data)
+        if self.result:
+            task_data.result = self.result
+        task_custom_data = CustomData(
+            type=self.name,
+            data=task_data.model_dump(),
+        )
+        if writer:
+            task_custom_data.dispatch(writer)
+        return task_custom_data.to_langchain()
+
+    def start(self, writer: StreamWriter | None = None, data: dict = {}) -> BaseMessage:
+        self.state = "new"
+        task_message = self._generate_and_dispatch_message(writer, data)
+        return task_message
+
+    def write_data(self, writer: StreamWriter | None = None, data: dict = {}) -> BaseMessage:
+        if self.state == "complete":
+            raise ValueError("Only incomplete tasks can output data.")
+        self.state = "running"
+        task_message = self._generate_and_dispatch_message(writer, data)
+        return task_message
+
+    def finish(
+        self,
+        result: Literal["success", "error"],
+        writer: StreamWriter | None = None,
+        data: dict = {},
+    ) -> BaseMessage:
+        self.state = "complete"
+        self.result = result
+        task_message = self._generate_and_dispatch_message(writer, data)
+        return task_message

src/agents/chatbot.py

@@ -0,0 +1,23 @@
+from langchain_core.messages import BaseMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.func import entrypoint
+
+from core import get_model, settings
+
+
+@entrypoint()
+async def chatbot(
+    inputs: dict[str, list[BaseMessage]],
+    *,
+    previous: dict[str, list[BaseMessage]],
+    config: RunnableConfig,
+):
+    messages = inputs["messages"]
+    if previous:
+        messages = previous["messages"] + messages
+
+    model = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    response = await model.ainvoke(messages)
+    return entrypoint.final(
+        value={"messages": [response]}, save={"messages": messages + [response]}
+    )

src/agents/command_agent.py

@@ -0,0 +1,55 @@
+import random
+from typing import Literal
+
+from langchain_core.messages import AIMessage
+from langgraph.graph import START, MessagesState, StateGraph
+from langgraph.types import Command
+
+
+class AgentState(MessagesState, total=False):
+    """`total=False` is PEP589 specs.
+
+    documentation: https://typing.readthedocs.io/en/latest/spec/typeddict.html#totality
+    """
+
+
+# Define the nodes
+
+
+def node_a(state: AgentState) -> Command[Literal["node_b", "node_c"]]:
+    print("Called A")
+    value = random.choice(["a", "b"])
+    goto: Literal["node_b", "node_c"]
+    # this is a replacement for a conditional edge function
+    if value == "a":
+        goto = "node_b"
+    else:
+        goto = "node_c"
+
+    # note how Command allows you to BOTH update the graph state AND route to the next node
+    return Command(
+        # this is the state update
+        update={"messages": [AIMessage(content=f"Hello {value}")]},
+        # this is a replacement for an edge
+        goto=goto,
+    )
+
+
+def node_b(state: AgentState):
+    print("Called B")
+    return {"messages": [AIMessage(content="Hello B")]}
+
+
+def node_c(state: AgentState):
+    print("Called C")
+    return {"messages": [AIMessage(content="Hello C")]}
+
+
+builder = StateGraph(AgentState)
+builder.add_edge(START, "node_a")
+builder.add_node(node_a)
+builder.add_node(node_b)
+builder.add_node(node_c)
+# NOTE: there are no edges between nodes A, B and C!
+
+command_agent = builder.compile()

src/agents/github_mcp_agent/github_mcp_agent.py

@@ -0,0 +1,102 @@
+"""GitHub MCP Agent - An agent that uses GitHub MCP tools for repository management."""
+
+import logging
+from datetime import datetime
+
+from langchain.agents import create_agent
+from langchain_core.tools import BaseTool
+from langchain_mcp_adapters.client import MultiServerMCPClient
+from langchain_mcp_adapters.sessions import StreamableHttpConnection
+from langgraph.graph.state import CompiledStateGraph
+
+from agents.lazy_agent import LazyLoadingAgent
+from core import get_model, settings
+
+logger = logging.getLogger(__name__)
+
+current_date = datetime.now().strftime("%B %d, %Y")
+prompt = f"""
+You are GitHubBot, a specialized assistant for GitHub repository management and development workflows.
+You have access to GitHub MCP tools that allow you to interact with GitHub repositories, issues, pull requests,
+and other GitHub resources. Today's date is {current_date}.
+
+Your capabilities include:
+- Repository management (create, clone, browse)
+- Issue management (create, list, update, close)
+- Pull request management (create, review, merge)
+- Branch management (create, switch, merge)
+- File operations (read, write, search)
+- Commit operations (create, view history)
+
+Guidelines:
+- Always be helpful and provide clear explanations of GitHub operations
+- When creating or modifying content, ensure it follows best practices
+- Be cautious with destructive operations (deletes, force pushes, etc.)
+- Provide context about what you're doing and why
+- Use appropriate commit messages and PR descriptions
+- Respect repository permissions and access controls
+
+NOTE: You have access to GitHub MCP tools that provide direct GitHub API access.
+"""
+
+
+class GitHubMCPAgent(LazyLoadingAgent):
+    """GitHub MCP Agent with async initialization."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._mcp_tools: list[BaseTool] = []
+        self._mcp_client: MultiServerMCPClient | None = None
+
+    async def load(self) -> None:
+        """Initialize the GitHub MCP agent by loading MCP tools."""
+        if not settings.GITHUB_PAT:
+            logger.info("GITHUB_PAT is not set, GitHub MCP agent will have no tools")
+            self._mcp_tools = []
+            self._graph = self._create_graph()
+            self._loaded = True
+            return
+
+        try:
+            # Initialize MCP client directly
+            github_pat = settings.GITHUB_PAT.get_secret_value()
+            connections = {
+                "github": StreamableHttpConnection(
+                    transport="streamable_http",
+                    url=settings.MCP_GITHUB_SERVER_URL,
+                    headers={
+                        "Authorization": f"Bearer {github_pat}",
+                    },
+                )
+            }
+
+            self._mcp_client = MultiServerMCPClient(connections)
+            logger.info("MCP client initialized successfully")
+
+            # Get tools from the client
+            self._mcp_tools = await self._mcp_client.get_tools()
+            logger.info(f"GitHub MCP agent initialized with {len(self._mcp_tools)} tools")
+
+        except Exception as e:
+            logger.error(f"Failed to initialize GitHub MCP agent: {e}")
+            self._mcp_tools = []
+            self._mcp_client = None
+
+        # Create and store the graph
+        self._graph = self._create_graph()
+        self._loaded = True
+
+    def _create_graph(self) -> CompiledStateGraph:
+        """Create the GitHub MCP agent graph."""
+        model = get_model(settings.DEFAULT_MODEL)
+
+        return create_agent(
+            model=model,
+            tools=self._mcp_tools,
+            name="github-mcp-agent",
+            system_prompt=prompt,
+        )
+
+
+# Create the agent instance
+github_mcp_agent = GitHubMCPAgent()

src/agents/interrupt_agent.py

@@ -0,0 +1,232 @@
+import logging
+from datetime import datetime
+from typing import Any
+
+from langchain_core.language_models.base import LanguageModelInput
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+from langchain_core.prompts import SystemMessagePromptTemplate
+from langchain_core.runnables import Runnable, RunnableConfig, RunnableLambda, RunnableSerializable
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.store.base import BaseStore
+from langgraph.types import interrupt
+from pydantic import BaseModel, Field
+
+from core import get_model, settings
+
+# Added logger
+logger = logging.getLogger(__name__)
+
+
+class AgentState(MessagesState, total=False):
+    """`total=False` is PEP589 specs.
+
+    documentation: https://typing.readthedocs.io/en/latest/spec/typeddict.html#totality
+    """
+
+    birthdate: datetime | None
+
+
+def wrap_model(
+    model: BaseChatModel | Runnable[LanguageModelInput, Any], system_prompt: BaseMessage
+) -> RunnableSerializable[AgentState, Any]:
+    preprocessor = RunnableLambda(
+        lambda state: [system_prompt] + state["messages"],
+        name="StateModifier",
+    )
+    return preprocessor | model
+
+
+background_prompt = SystemMessagePromptTemplate.from_template("""
+You are a helpful assistant that tells users there zodiac sign.
+Provide a one sentence summary of the origin of zodiac signs.
+Don't tell the user what their sign is, you are just demonstrating your knowledge on the topic.
+""")
+
+
+async def background(state: AgentState, config: RunnableConfig) -> AgentState:
+    """This node is to demonstrate doing work before the interrupt"""
+
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(m, background_prompt.format())
+    response = await model_runnable.ainvoke(state, config)
+
+    return {"messages": [AIMessage(content=response.content)]}
+
+
+birthdate_extraction_prompt = SystemMessagePromptTemplate.from_template("""
+You are an expert at extracting birthdates from conversational text.
+
+Rules for extraction:
+- Look for user messages that mention birthdates
+- Consider various date formats (MM/DD/YYYY, YYYY-MM-DD, Month Day, Year)
+- Validate that the date is reasonable (not in the future)
+- If no clear birthdate was provided by the user, return None
+""")
+
+
+class BirthdateExtraction(BaseModel):
+    birthdate: str | None = Field(
+        description="The extracted birthdate in YYYY-MM-DD format. If no birthdate is found, this should be None."
+    )
+    reasoning: str = Field(
+        description="Explanation of how the birthdate was extracted or why no birthdate was found"
+    )
+
+
+async def determine_birthdate(
+    state: AgentState, config: RunnableConfig, store: BaseStore
+) -> AgentState:
+    """This node examines the conversation history to determine user's birthdate, checking store first."""
+
+    # Attempt to get user_id for unique storage per user
+    user_id = config["configurable"].get("user_id")
+    logger.info(f"[determine_birthdate] Extracted user_id: {user_id}")
+    namespace = None
+    key = "birthdate"
+    birthdate = None  # Initialize birthdate
+
+    if user_id:
+        # Use user_id in the namespace to ensure uniqueness per user
+        namespace = (user_id,)
+
+        # Check if we already have the birthdate in the store for this user
+        try:
+            result = await store.aget(namespace, key=key)
+            # Handle cases where store.aget might return Item directly or a list
+            user_data = None
+            if result:  # Check if anything was returned
+                if isinstance(result, list):
+                    if result:  # Check if list is not empty
+                        user_data = result[0]
+                else:  # Assume it's the Item object directly
+                    user_data = result
+
+            if user_data and user_data.value.get("birthdate"):
+                # Convert ISO format string back to datetime object
+                birthdate_str = user_data.value["birthdate"]
+                birthdate = datetime.fromisoformat(birthdate_str) if birthdate_str else None
+                # We already have the birthdate, return it
+                logger.info(
+                    f"[determine_birthdate] Found birthdate in store for user {user_id}: {birthdate}"
+                )
+                return {
+                    "birthdate": birthdate,
+                    "messages": [],
+                }
+        except Exception as e:
+            # Log the error or handle cases where the store might be unavailable
+            logger.error(f"Error reading from store for namespace {namespace}, key {key}: {e}")
+            # Proceed with extraction if read fails
+            pass
+    else:
+        # If no user_id, we cannot reliably store/retrieve user-specific data.
+        # Consider logging this situation.
+        logger.warning(
+            "Warning: user_id not found in config. Skipping persistent birthdate storage/retrieval for this run."
+        )
+
+    # If birthdate wasn't retrieved from store, proceed with extraction
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(
+        m.with_structured_output(BirthdateExtraction), birthdate_extraction_prompt.format()
+    ).with_config(tags=["skip_stream"])
+    response: BirthdateExtraction = await model_runnable.ainvoke(state, config)
+
+    # If no birthdate found after extraction attempt, interrupt
+    if response.birthdate is None:
+        birthdate_input = interrupt(f"{response.reasoning}\nPlease tell me your birthdate?")
+        # Re-run extraction with the new input
+        state["messages"].append(HumanMessage(birthdate_input))
+        # Note: Recursive call might need careful handling of depth or state updates
+        return await determine_birthdate(state, config, store)
+
+    # Birthdate found - convert string to datetime
+    try:
+        birthdate = datetime.fromisoformat(response.birthdate)
+    except ValueError:
+        # If parsing fails, ask for clarification
+        birthdate_input = interrupt(
+            "I couldn't understand the date format. Please provide your birthdate in YYYY-MM-DD format."
+        )
+        # Re-run extraction with the new input
+        state["messages"].append(HumanMessage(birthdate_input))
+        # Note: Recursive call might need careful handling of depth or state updates
+        return await determine_birthdate(state, config, store)
+
+    # Store the newly extracted birthdate only if we have a user_id
+    if user_id and namespace:
+        # Convert datetime to ISO format string for JSON serialization
+        birthdate_str = birthdate.isoformat() if birthdate else None
+        try:
+            await store.aput(namespace, key, {"birthdate": birthdate_str})
+        except Exception as e:
+            # Log the error or handle cases where the store write might fail
+            logger.error(f"Error writing to store for namespace {namespace}, key {key}: {e}")
+
+    # Return the determined birthdate (either from store or extracted)
+    logger.info(f"[determine_birthdate] Returning birthdate {birthdate} for user {user_id}")
+    return {
+        "birthdate": birthdate,
+        "messages": [],
+    }
+
+
+response_prompt = SystemMessagePromptTemplate.from_template("""
+You are a helpful assistant.
+
+Known information:
+- The user's birthdate is {birthdate_str}
+
+User's latest message: "{last_user_message}"
+
+Based on the known information and the user's message, provide a helpful and relevant response.
+If the user asked for their birthdate, confirm it.
+If the user asked for their zodiac sign, calculate it and tell them.
+Otherwise, respond conversationally based on their message.
+""")
+
+
+async def generate_response(state: AgentState, config: RunnableConfig) -> AgentState:
+    """Generates the final response based on the user's query and the available birthdate."""
+    birthdate = state.get("birthdate")
+    if state.get("messages") and isinstance(state["messages"][-1], HumanMessage):
+        last_user_message = state["messages"][-1].content
+    else:
+        last_user_message = ""
+
+    if not birthdate:
+        # This should ideally not be reached if determine_birthdate worked correctly and possibly interrupted.
+        # Handle cases where birthdate might still be missing.
+        return {
+            "messages": [
+                AIMessage(
+                    content="I couldn't determine your birthdate. Could you please provide it?"
+                )
+            ]
+        }
+
+    birthdate_str = birthdate.strftime("%B %d, %Y")  # Format for display
+
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(
+        m, response_prompt.format(birthdate_str=birthdate_str, last_user_message=last_user_message)
+    )
+    response = await model_runnable.ainvoke(state, config)
+
+    return {"messages": [AIMessage(content=response.content)]}
+
+
+# Define the graph
+agent = StateGraph(AgentState)
+agent.add_node("background", background)
+agent.add_node("determine_birthdate", determine_birthdate)
+agent.add_node("generate_response", generate_response)
+
+agent.set_entry_point("background")
+agent.add_edge("background", "determine_birthdate")
+agent.add_edge("determine_birthdate", "generate_response")
+agent.add_edge("generate_response", END)
+
+interrupt_agent = agent.compile()
+interrupt_agent.name = "interrupt-agent"

src/agents/knowledge_base_agent.py

@@ -0,0 +1,174 @@
+import logging
+import os
+from typing import Any
+
+from langchain_aws import AmazonKnowledgeBasesRetriever
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnableSerializable
+from langchain_core.runnables.base import RunnableSequence
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.managed import RemainingSteps
+
+from core import get_model, settings
+
+logger = logging.getLogger(__name__)
+
+
+# Define the state
+class AgentState(MessagesState, total=False):
+    """State for Knowledge Base agent."""
+
+    remaining_steps: RemainingSteps
+    retrieved_documents: list[dict[str, Any]]
+    kb_documents: str
+
+
+# Create the retriever
+def get_kb_retriever():
+    """Create and return a Knowledge Base retriever instance."""
+    # Get the Knowledge Base ID from environment
+    kb_id = os.environ.get("AWS_KB_ID", "")
+    if not kb_id:
+        raise ValueError("AWS_KB_ID environment variable must be set")
+
+    # Create the retriever with the specified Knowledge Base ID
+    retriever = AmazonKnowledgeBasesRetriever(
+        knowledge_base_id=kb_id,
+        retrieval_config={
+            "vectorSearchConfiguration": {
+                "numberOfResults": 3,
+            }
+        },
+    )
+    return retriever
+
+
+def wrap_model(model: BaseChatModel) -> RunnableSerializable[AgentState, AIMessage]:
+    """Wrap the model with a system prompt for the Knowledge Base agent."""
+
+    def create_system_message(state):
+        base_prompt = """You are a helpful assistant that provides accurate information based on retrieved documents.
+
+        You will receive a query along with relevant documents retrieved from a knowledge base. Use these documents to inform your response.
+
+        Follow these guidelines:
+        1. Base your answer primarily on the retrieved documents
+        2. If the documents contain the answer, provide it clearly and concisely
+        3. If the documents are insufficient, state that you don't have enough information
+        4. Never make up facts or information not present in the documents
+        5. Always cite the source documents when referring to specific information
+        6. If the documents contradict each other, acknowledge this and explain the different perspectives
+
+        Format your response in a clear, conversational manner. Use markdown formatting when appropriate.
+        """
+
+        # Check if documents were retrieved
+        if "kb_documents" in state:
+            # Append document information to the system prompt
+            document_prompt = f"\n\nI've retrieved the following documents that may be relevant to the query:\n\n{state['kb_documents']}\n\nPlease use these documents to inform your response to the user's query. Only use information from these documents and clearly indicate when you are unsure."
+            return [SystemMessage(content=base_prompt + document_prompt)] + state["messages"]
+        else:
+            # No documents were retrieved
+            no_docs_prompt = (
+                "\n\nNo relevant documents were found in the knowledge base for this query."
+            )
+            return [SystemMessage(content=base_prompt + no_docs_prompt)] + state["messages"]
+
+    preprocessor = RunnableLambda(
+        create_system_message,
+        name="StateModifier",
+    )
+    return RunnableSequence(preprocessor, model)
+
+
+async def retrieve_documents(state: AgentState, config: RunnableConfig) -> AgentState:
+    """Retrieve relevant documents from the knowledge base."""
+    # Get the last human message
+    human_messages = [msg for msg in state["messages"] if isinstance(msg, HumanMessage)]
+    if not human_messages:
+        # Include messages from original state
+        return {"messages": [], "retrieved_documents": []}
+
+    # Use the last human message as the query
+    query = human_messages[-1].content
+
+    try:
+        # Initialize the retriever
+        retriever = get_kb_retriever()
+
+        # Retrieve documents
+        retrieved_docs = await retriever.ainvoke(query)
+
+        # Create document summaries for the state
+        document_summaries = []
+        for i, doc in enumerate(retrieved_docs, 1):
+            summary = {
+                "id": doc.metadata.get("id", f"doc-{i}"),
+                "source": doc.metadata.get("source", "Unknown"),
+                "title": doc.metadata.get("title", f"Document {i}"),
+                "content": doc.page_content,
+                "relevance_score": doc.metadata.get("score", 0),
+            }
+            document_summaries.append(summary)
+
+        logger.info(f"Retrieved {len(document_summaries)} documents for query: {query[:50]}...")
+
+        return {"retrieved_documents": document_summaries, "messages": []}
+
+    except Exception as e:
+        logger.error(f"Error retrieving documents: {str(e)}")
+        return {"retrieved_documents": [], "messages": []}
+
+
+async def prepare_augmented_prompt(state: AgentState, config: RunnableConfig) -> AgentState:
+    """Prepare a prompt augmented with retrieved document content."""
+    # Get retrieved documents
+    documents = state.get("retrieved_documents", [])
+
+    if not documents:
+        return {"messages": []}
+
+    # Format retrieved documents for the model
+    formatted_docs = "\n\n".join(
+        [
+            f"--- Document {i + 1} ---\n"
+            f"Source: {doc.get('source', 'Unknown')}\n"
+            f"Title: {doc.get('title', 'Unknown')}\n\n"
+            f"{doc.get('content', '')}"
+            for i, doc in enumerate(documents)
+        ]
+    )
+
+    # Store formatted documents in the state
+    return {"kb_documents": formatted_docs, "messages": []}
+
+
+async def acall_model(state: AgentState, config: RunnableConfig) -> AgentState:
+    """Generate a response based on the retrieved documents."""
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(m)
+
+    response = await model_runnable.ainvoke(state, config)
+
+    return {"messages": [response]}
+
+
+# Define the graph
+agent = StateGraph(AgentState)
+
+# Add nodes
+agent.add_node("retrieve_documents", retrieve_documents)
+agent.add_node("prepare_augmented_prompt", prepare_augmented_prompt)
+agent.add_node("model", acall_model)
+
+# Set entry point
+agent.set_entry_point("retrieve_documents")
+
+# Add edges to define the flow
+agent.add_edge("retrieve_documents", "prepare_augmented_prompt")
+agent.add_edge("prepare_augmented_prompt", "model")
+agent.add_edge("model", END)
+
+# Compile the agent
+kb_agent = agent.compile()

src/agents/langgraph_supervisor_agent.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from langchain.agents import create_agent
+from langgraph_supervisor import create_supervisor
+
+from core import get_model, settings
+
+model = get_model(settings.DEFAULT_MODEL)
+
+
+def add(a: float, b: float) -> float:
+    """Add two numbers."""
+    return a + b
+
+
+def multiply(a: float, b: float) -> float:
+    """Multiply two numbers."""
+    return a * b
+
+
+def web_search(query: str) -> str:
+    """Search the web for information."""
+    return (
+        "Here are the headcounts for each of the FAANG companies in 2024:\n"
+        "1. **Facebook (Meta)**: 67,317 employees.\n"
+        "2. **Apple**: 164,000 employees.\n"
+        "3. **Amazon**: 1,551,000 employees.\n"
+        "4. **Netflix**: 14,000 employees.\n"
+        "5. **Google (Alphabet)**: 181,269 employees."
+    )
+
+
+math_agent: Any = create_agent(
+    model=model,
+    tools=[add, multiply],
+    name="sub-agent-math_expert",
+    system_prompt="You are a math expert. Always use one tool at a time.",
+).with_config(tags=["skip_stream"])
+
+research_agent: Any = create_agent(
+    model=model,
+    tools=[web_search],
+    name="sub-agent-research_expert",
+    system_prompt="You are a world class researcher with access to web search. Do not do any math.",
+).with_config(tags=["skip_stream"])
+
+
+# Create supervisor workflow
+workflow = create_supervisor(
+    [research_agent, math_agent],
+    model=model,
+    prompt=(
+        "You are a team supervisor managing a research expert and a math expert. "
+        "For current events, use research_agent. "
+        "For math problems, use math_agent."
+    ),
+    add_handoff_back_messages=True,
+    # UI now expects this to be True so we don't have to guess when a handoff back occurs
+    output_mode="full_history",  # otherwise when reloading conversations, the sub-agents' messages are not included
+)
+
+langgraph_supervisor_agent = workflow.compile()

src/agents/langgraph_supervisor_hierarchy_agent.py

@@ -0,0 +1,46 @@
+from langchain.agents import create_agent
+from langgraph_supervisor import create_supervisor
+
+from agents.langgraph_supervisor_agent import add, multiply, web_search
+from core import get_model, settings
+
+model = get_model(settings.DEFAULT_MODEL)
+
+
+def workflow(chosen_model):
+    math_agent = create_agent(
+        model=chosen_model,
+        tools=[add, multiply],
+        name="sub-agent-math_expert",  # Identify the graph node as a sub-agent
+        system_prompt="You are a math expert. Always use one tool at a time.",
+    ).with_config(tags=["skip_stream"])
+
+    research_agent = (
+        create_supervisor(
+            [math_agent],
+            model=chosen_model,
+            tools=[web_search],
+            prompt="You are a world class researcher with access to web search. Do not do any math, you have a math expert for that. ",
+            supervisor_name="supervisor-research_expert",  # Identify the graph node as a supervisor to the math agent
+        )
+        .compile(
+            name="sub-agent-research_expert"
+        )  # Identify the graph node as a sub-agent to the main supervisor
+        .with_config(tags=["skip_stream"])
+    )  # Stream tokens are ignored for sub-agents in the UI
+
+    # Create supervisor workflow
+    return create_supervisor(
+        [research_agent],
+        model=chosen_model,
+        prompt=(
+            "You are a team supervisor managing a research expert with math capabilities."
+            "For current events, use research_agent. "
+        ),
+        add_handoff_back_messages=True,
+        # UI now expects this to be True so we don't have to guess when a handoff back occurs
+        output_mode="full_history",  # otherwise when reloading conversations, the sub-agents' messages are not included
+    )  # default name for supervisor is "supervisor".
+
+
+langgraph_supervisor_hierarchy_agent = workflow(model).compile()

src/agents/lazy_agent.py

@@ -0,0 +1,43 @@
+"""Agent types with async initialization and dynamic graph creation."""
+
+from abc import ABC, abstractmethod
+
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.pregel import Pregel
+
+
+class LazyLoadingAgent(ABC):
+    """Base class for agents that require async loading."""
+
+    def __init__(self) -> None:
+        """Initialize the agent."""
+        self._loaded = False
+        self._graph: CompiledStateGraph | Pregel | None = None
+
+    @abstractmethod
+    async def load(self) -> None:
+        """
+        Perform async loading for this agent.
+
+        This method is called during service startup and should handle:
+        - Setting up external connections (MCP clients, databases, etc.)
+        - Loading tools or resources
+        - Any other async setup required
+        - Creating the agent's graph
+        """
+        raise NotImplementedError  # pragma: no cover
+
+    def get_graph(self) -> CompiledStateGraph | Pregel:
+        """
+        Get the agent's graph.
+
+        Returns the graph instance that was created during load().
+
+        Returns:
+            The agent's graph (CompiledStateGraph or Pregel)
+        """
+        if not self._loaded:
+            raise RuntimeError("Agent not loaded. Call load() first.")
+        if self._graph is None:
+            raise RuntimeError("Agent graph not created during load().")
+        return self._graph

src/agents/llama_guard.py

@@ -0,0 +1,121 @@
+from enum import Enum
+
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from langchain_core.prompts import PromptTemplate
+from pydantic import BaseModel, Field
+
+from core import get_model, settings
+from schema.models import GroqModelName
+
+
+class SafetyAssessment(Enum):
+    SAFE = "safe"
+    UNSAFE = "unsafe"
+    ERROR = "error"
+
+
+class LlamaGuardOutput(BaseModel):
+    safety_assessment: SafetyAssessment = Field(description="The safety assessment of the content.")
+    unsafe_categories: list[str] = Field(
+        description="If content is unsafe, the list of unsafe categories.", default=[]
+    )
+
+
+unsafe_content_categories = {
+    "S1": "Violent Crimes.",
+    "S2": "Non-Violent Crimes.",
+    "S3": "Sex Crimes.",
+    "S4": "Child Exploitation.",
+    "S5": "Defamation.",
+    "S6": "Specialized Advice.",
+    "S7": "Privacy.",
+    "S8": "Intellectual Property.",
+    "S9": "Indiscriminate Weapons.",
+    "S10": "Hate.",
+    "S11": "Self-Harm.",
+    "S12": "Sexual Content.",
+    "S13": "Elections.",
+    "S14": "Code Interpreter Abuse.",
+}
+
+categories_str = "\n".join([f"{k}: {v}" for k, v in unsafe_content_categories.items()])
+
+llama_guard_instructions = f"""
+Task: Check if there is unsafe content in '{{role}}' messages in conversations according to our safety policy with the below categories.
+
+<BEGIN UNSAFE CONTENT CATEGORIES>
+{categories_str}
+<END UNSAFE CONTENT CATEGORIES>
+
+<BEGIN CONVERSATION>
+
+{{conversation_history}}
+
+<END CONVERSATION>
+
+Provide your safety assessment for ONLY THE LAST {{role}} in the above conversation:
+- First line must read 'safe' or 'unsafe'.
+- If unsafe, a second line must include a comma-separated list of violated categories."""
+
+
+def parse_llama_guard_output(output: str) -> LlamaGuardOutput:
+    if output == "safe":
+        return LlamaGuardOutput(safety_assessment=SafetyAssessment.SAFE)
+    parsed_output = output.split("\n")
+    if len(parsed_output) != 2 or parsed_output[0] != "unsafe":
+        return LlamaGuardOutput(safety_assessment=SafetyAssessment.ERROR)
+    try:
+        categories = parsed_output[1].split(",")
+        readable_categories = [unsafe_content_categories[c.strip()].strip(".") for c in categories]
+        return LlamaGuardOutput(
+            safety_assessment=SafetyAssessment.UNSAFE,
+            unsafe_categories=readable_categories,
+        )
+    except KeyError:
+        return LlamaGuardOutput(safety_assessment=SafetyAssessment.ERROR)
+
+
+class LlamaGuard:
+    def __init__(self) -> None:
+        if settings.GROQ_API_KEY is None:
+            print("GROQ_API_KEY not set, skipping LlamaGuard")
+            self.model = None
+            return
+        self.model = get_model(GroqModelName.LLAMA_GUARD_4_12B).with_config(tags=["skip_stream"])
+        self.prompt = PromptTemplate.from_template(llama_guard_instructions)
+
+    def _compile_prompt(self, role: str, messages: list[AnyMessage]) -> str:
+        role_mapping = {"ai": "Agent", "human": "User"}
+        messages_str = [
+            f"{role_mapping[m.type]}: {m.content}" for m in messages if m.type in ["ai", "human"]
+        ]
+        conversation_history = "\n\n".join(messages_str)
+        return self.prompt.format(role=role, conversation_history=conversation_history)
+
+    def invoke(self, role: str, messages: list[AnyMessage]) -> LlamaGuardOutput:
+        if self.model is None:
+            return LlamaGuardOutput(safety_assessment=SafetyAssessment.SAFE)
+        compiled_prompt = self._compile_prompt(role, messages)
+        result = self.model.invoke([HumanMessage(content=compiled_prompt)])
+        return parse_llama_guard_output(str(result.content))
+
+    async def ainvoke(self, role: str, messages: list[AnyMessage]) -> LlamaGuardOutput:
+        if self.model is None:
+            return LlamaGuardOutput(safety_assessment=SafetyAssessment.SAFE)
+        compiled_prompt = self._compile_prompt(role, messages)
+        result = await self.model.ainvoke([HumanMessage(content=compiled_prompt)])
+        return parse_llama_guard_output(str(result.content))
+
+
+if __name__ == "__main__":
+    llama_guard = LlamaGuard()
+    output = llama_guard.invoke(
+        "Agent",
+        [
+            HumanMessage(content="What's a good way to harm an animal?"),
+            AIMessage(
+                content="There are many ways to harm animals, but some include hitting them with a stick, throwing rocks at them, or poisoning them."
+            ),
+        ],
+    )
+    print(output)

src/agents/portfolio_agent/database_search.py

@@ -0,0 +1,44 @@
+from typing import List
+
+from langchain_core.documents import Document
+from langchain_core.tools import BaseTool, tool
+from memory.postgres import load_pgvector_retriever
+
+
+def hybrid_search(query: str, k: int = 6) -> List[Document]:
+    retriever = load_pgvector_retriever(k)
+    return retriever.invoke(query)
+
+
+def format_contexts(documents: List[Document]) -> str:
+    formatted_docs = []
+
+    for i, doc in enumerate(documents):
+        source = doc.metadata.get("source_db", "Unknown")
+        content = doc.page_content.strip()
+
+        formatted_docs.append(
+            f"--- Document {i + 1} (Source: {source}) ---\n{content}"
+        )
+
+    return "\n\n".join(formatted_docs)
+
+
+# DATABASE SEARCH TOOL
+def database_search_func(query: str) -> str:
+    """
+    Search Anuj Joshi's portfolio database for any information.
+    Use this tool whenever the user asks about education, experience, projects, blog posts, testimonials, or any other information NOT available 
+    in your system prompt.
+    """
+
+    documents = hybrid_search(query, k=6)
+
+    if not documents:
+        return "No relevant portfolio information found."
+
+    return format_contexts(documents)
+
+
+database_search: BaseTool = tool(database_search_func)
+database_search.name = "Database_Search"

src/agents/portfolio_agent/portfolio_agent.py

@@ -0,0 +1,85 @@
+from typing import Literal, List
+from langchain_core.messages import AIMessage, SystemMessage
+from langchain_core.runnables import RunnableConfig, RunnableLambda
+from langgraph.graph import END, MessagesState, StateGraph
+from agents.portfolio_agent.prompt import SYSTEM_PROMPT, FOLLOWUP_GENERATION_PROMPT
+from langgraph.prebuilt import ToolNode
+from pydantic import BaseModel, Field
+from core import get_model, settings
+from agents.portfolio_agent.database_search import database_search
+
+# State Extension
+class AgentState(MessagesState):
+    """Agent state for the portfolio assistant."""
+    summary: str
+    follow_up: List[str]
+
+class FollowUpQuestions(BaseModel):
+    questions: List[str] = Field(description="A list of follow-up questions")
+
+# Tools
+tools = [database_search]
+
+# Portfolio Model Wrapper (WITH tools)
+def wrap_portfolio_model(model):
+    model = model.bind_tools(tools)
+
+    def prepare_messages(state: AgentState):
+        summary = state.get("summary", "")
+        system_content = SYSTEM_PROMPT
+        if summary:
+            system_content += f"\n\n### Previous Conversation Summary:\n{summary}"
+        
+        return [SystemMessage(content=system_content)] + state["messages"]
+
+    return RunnableLambda(prepare_messages) | model
+
+async def call_portfolio_model(state: AgentState, config: RunnableConfig):
+    model = get_model(settings.DEFAULT_MODEL)
+    runnable = wrap_portfolio_model(model)
+    response = await runnable.ainvoke(state, config)
+    return {"messages": [response]}
+
+async def generate_followup(state: AgentState, config: RunnableConfig):
+    model = get_model(settings.DEFAULT_MODEL)
+    chain = model.with_structured_output(FollowUpQuestions)
+    print("Followup messages: ", state["messages"])
+    messages = [SystemMessage(content=FOLLOWUP_GENERATION_PROMPT)] + state["messages"]
+    
+    try:
+        response = await chain.ainvoke(messages, config)
+        return {"follow_up": response.questions}
+    except Exception:
+        return {"follow_up": ["What are his contact details?", "What are his projects?", "What are his skills?", "What are his work experiences?", "What are his achievements?", "What are his leadership roles?"]}
+
+
+# Tool Routing
+def route_after_model(state: AgentState) -> Literal["tools", "generate_followup"]:
+    last_message = state["messages"][-1]
+    
+    if isinstance(last_message, AIMessage) and last_message.tool_calls:
+        return "tools"
+    return "generate_followup"
+
+# Graph Definition
+graph = StateGraph(AgentState)
+
+graph.add_node("agent", call_portfolio_model)
+graph.add_node("tools", ToolNode(tools))
+graph.add_node("generate_followup", generate_followup)
+
+graph.set_entry_point("agent")
+
+graph.add_conditional_edges(
+    "agent",
+    route_after_model,
+    {
+        "tools": "tools",
+        "generate_followup": "generate_followup",
+    },
+)
+
+graph.add_edge("tools", "agent")
+graph.add_edge("generate_followup", END)
+
+portfolio_agent = graph.compile(debug=True)

src/agents/portfolio_agent/prompt.py

@@ -0,0 +1,115 @@
+PORTFOLIO_URL = "https://anujjoshi.netlify.app"
+
+SYSTEM_PROMPT = f"""
+You are an award-winning Professional Portfolio Assistant representing Anuj Joshi.
+Your goal is to answer questions from visitors about Anuj's skills, projects, and experience and also cite the sources of information.
+You are NOT chatting with Anuj Joshi himself, but with a recruiter, potential employer, or visitor of his portfolio.
+
+Name: Anuj Joshi  
+Role: Full Stack Developer | AI & Machine Learning Engineer  
+Location: New Delhi, India  
+
+# CONTACT ({PORTFOLIO_URL}/contact)
+- Email: anujjoshi3105@gmail.com
+- LinkedIn: https://www.linkedin.com/in/anujjoshi3105/
+- X (Twitter): https://x.com/anujjoshi3105
+
+# Competitive Programming ({PORTFOLIO_URL}/about)
+- LeetCode: https://leetcode.com/u/anujjoshi3105 Max Rating: 1910, Level: Knight, 750+ Problems Solved 
+- Codeforces: https://codeforces.com/profile/anujjoshi3105 Max Rating: 1434, Level: specialist
+- AtCoder: https://atcoder.jp/users/anujjoshi3105 Max Rating: 929, Level: Green
+- GeeksforGeeks: https://www.geeksforgeeks.org/profile/anujjoshi3105 Institute Rank: 46
+
+# Summary ({PORTFOLIO_URL}/about)
+Anuj Joshi is a full-stack developer and AI engineer with hands-on experience building
+production-grade AI agents, healthcare assistants, computer vision systems, and scalable
+web platforms. His work spans AI agents, LLM systems, backend engineering, and applied ML,
+with experience in startups, research-driven teams, and academic organizations.
+
+# EDUCATION ({PORTFOLIO_URL}/about#education)
+## Bachelor of Technology (B.Tech)
+- Field: Computer Science & Engineering
+- Minor: Machine Learning
+- Institution: Delhi Technological University (DTU), Delhi, India
+- Duration: November 2022 – May 2026
+- CGPA: 9.35 / 10
+- Strong academic performance in Machine Learning, Deep Learning, and Artificial Intelligence
+- Core coursework includes: Machine Learning, Deep Learning, Artificial Intelligence, Operating Systems, Database Management Systems, Computer Networks
+
+## CBSE Class XII
+https://drive.google.com/file/d/14EcEdGaikR0dynanY7NGLPGMNMUhx9Ds/view?usp=sharing
+- Stream: PCM (Physics, Chemistry, Mathematics)
+- Institution: Vivekanand International School, Delhi, India
+- Duration: April 2021 – July 2022
+- Score: 98.8%
+- Perfect scores in Mathematics and Chemistry
+
+## CBSE Class X
+https://drive.google.com/file/d/14CHsmHp3kvbjze9o3cMxxihrKo5IMRyI/view?usp=sharing
+- Institution: Vivekanand International School, Delhi, India
+- Duration: April 2019 – March 2020
+- Score: 97.0%
+- Scored 99 in Mathematics, Science, and Computer Science
+
+# WORK EXPERIENCE ({PORTFOLIO_URL}/about#experience)
+## 1) Full Stack Developer Intern – Quickintell (Remote, 2025)
+- Built secure AI voice and chat agents integrated with EHR APIs for HIPAA-compliant patient identity verification.
+- Developed production-ready healthcare assistants using LangChain, AWS Lambda, and vector databases.
+- Improved retrieval performance and system scalability.
+
+## 2) Software Developer Intern – ITP Electronics (Delhi, 2024)
+https://drive.google.com/file/d/15Jzu-oujhKUDZiWoQf1WWWN6ZJnmBqQH/view?usp=drive_link
+- Developed computer vision solutions for automated wire harness detection.
+- Reduced manual inspection time using OpenCV-based vision pipelines.
+- Built image-to-BoM automation systems using backend services.
+
+## 3) Web Developer Intern – USIP-DTU (Delhi, 2024)
+https://drive.google.com/file/d/150EAtBVjP1DV-b_v0JKhVYzhIVoCvAWO/view
+- Built a full-stack ERP system using PHP, Node.js, and MySQL.
+- Implemented JWT-based role-based access control (RBAC).
+- Automated proposal lifecycle and document management for 50+ users.
+
+# PROJECTS ({PORTFOLIO_URL}/project#projects)
+- Ekalavya – AI-powered EdTech SaaS with adaptive learning, quizzes, and AI tutors.
+- BITLOG – Developer-focused blogging platform with authentication and scalability.
+- Industrial Research & Development Centre Portal – Research workflow automation platform.
+- NicoGauge – ML-based evaluation platform for learning analytics.
+- Fictiora – Entertainment discovery platform using Next.js and TMDB APIs.
+For more info visit {PORTFOLIO_URL}/project
+
+# TECHNICAL SKILLS ({PORTFOLIO_URL}/about#skills)
+- AI / ML: PyTorch, TensorFlow, Keras, Scikit-learn, CNNs, GANs, Transformers, LangChain, LLM-based agents
+- Backend / Full Stack: Python, Node.js, PHP, FastAPI, Flask, Django, Express.js, REST APIs, JWT Authentication, RBAC
+- Databases: PostgreSQL, MySQL, MongoDB, Vector Databases
+- Frontend: React, Next.js, TailwindCSS, JavaScript
+- DevOps / Tools: AWS Lambda, Docker, Git
+
+# LEADERSHIP ROLES & VOLUNTEERING ({PORTFOLIO_URL}/about#experience)
+- General Secretary, Society of Robotics (DTU): Led a 50+ member team and organized robotics events with 200+ participants.
+- Volunteer, Summer School on AI (DTU) https://drive.google.com/file/d/10Jx3yC8gmFYHkl0KXucaUOZJqtf9QkJq/view?usp=drive_link: Supported hands-on sessions on deep learning, transformers, and generative AI.
+
+# For other relevant information always use **Database_Search** tool and cite {PORTFOLIO_URL}/blog**
+
+# TOOL USAGE RULES (STRICT & CRITICAL):
+1. **INCOMPLETE INFORMATION**: This system prompt ONLY contains a basic overview. It does NOT contain blog posts, specific contest results, latest activities, or deep technical details.
+2. **SEARCH FIRST**: If a user asks about Contests/Competitions, Blog Posts/Articles or Technical Deep-Dives, You **MUST** call the `Database_Search` tool immediately.
+3. **NO "I DON'T KNOW" WITHOUT SEARCH**: Never tell the user information is unavailable until AFTER you have used `Database_Search`.
+4. **DIRECT ACTION**: Do NOT ask for permission to search. Do NOT explain that you are searching. Just call the tool.
+5. **CITE SOURCES**: Always mention that more info is available at {PORTFOLIO_URL}/blog?q=[relevant-query] and use the content from the tool accurately.
+
+# STYLE: Professional, concise, witty and helpful.
+"""
+
+FOLLOWUP_GENERATION_PROMPT = """
+Generate 3–5 unique short follow-up options for the user based on the last AI message.
+
+Rules:
+1. If the AI asked a question or offered a choice (e.g., to search the database), include relevant replies like "Accept", "Reject".
+2. If the AI provided information, generate follow-up questions to explore Anuj Joshi’s portfolio (skills, projects, experience, etc.).
+3. Keep options concise and unique (2-6 words).
+
+Return ONLY valid JSON:
+{
+  "questions": ["option 1", "option 2", ...]
+}
+"""

src/agents/rag_assistant.py

@@ -0,0 +1,146 @@
+from datetime import datetime
+from typing import Literal
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, SystemMessage
+from langchain_core.runnables import (
+    RunnableConfig,
+    RunnableLambda,
+    RunnableSerializable,
+)
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.managed import RemainingSteps
+from langgraph.prebuilt import ToolNode
+
+from agents.llama_guard import LlamaGuard, LlamaGuardOutput, SafetyAssessment
+from agents.tools import database_search
+from core import get_model, settings
+
+
+class AgentState(MessagesState, total=False):
+    """`total=False` is PEP589 specs.
+
+    documentation: https://typing.readthedocs.io/en/latest/spec/typeddict.html#totality
+    """
+
+    safety: LlamaGuardOutput
+    remaining_steps: RemainingSteps
+
+
+tools = [database_search]
+
+
+current_date = datetime.now().strftime("%B %d, %Y")
+instructions = f"""
+    You are AcmeBot, a helpful and knowledgeable virtual assistant designed to support employees by retrieving
+    and answering questions based on AcmeTech's official Employee Handbook. Your primary role is to provide
+    accurate, concise, and friendly information about company policies, values, procedures, and employee resources.
+    Today's date is {current_date}.
+
+    NOTE: THE USER CAN'T SEE THE TOOL RESPONSE.
+
+    A few things to remember:
+    - If you have access to multiple databases, gather information from a diverse range of sources before crafting your response.
+    - Please include markdown-formatted links to any citations used in your response. Only include one
+    or two citations per response unless more are needed. ONLY USE LINKS RETURNED BY THE TOOLS.
+    - Only use information from the database. Do not use information from outside sources.
+    """
+
+
+def wrap_model(model: BaseChatModel) -> RunnableSerializable[AgentState, AIMessage]:
+    bound_model = model.bind_tools(tools)
+    preprocessor = RunnableLambda(
+        lambda state: [SystemMessage(content=instructions)] + state["messages"],
+        name="StateModifier",
+    )
+    return preprocessor | bound_model  # type: ignore[return-value]
+
+
+def format_safety_message(safety: LlamaGuardOutput) -> AIMessage:
+    content = (
+        f"This conversation was flagged for unsafe content: {', '.join(safety.unsafe_categories)}"
+    )
+    return AIMessage(content=content)
+
+
+async def acall_model(state: AgentState, config: RunnableConfig) -> AgentState:
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(m)
+    response = await model_runnable.ainvoke(state, config)
+
+    # Run llama guard check here to avoid returning the message if it's unsafe
+    llama_guard = LlamaGuard()
+    safety_output = await llama_guard.ainvoke("Agent", state["messages"] + [response])
+    if safety_output.safety_assessment == SafetyAssessment.UNSAFE:
+        return {
+            "messages": [format_safety_message(safety_output)],
+            "safety": safety_output,
+        }
+
+    if state["remaining_steps"] < 2 and response.tool_calls:
+        return {
+            "messages": [
+                AIMessage(
+                    id=response.id,
+                    content="Sorry, need more steps to process this request.",
+                )
+            ]
+        }
+    # We return a list, because this will get added to the existing list
+    return {"messages": [response]}
+
+
+async def llama_guard_input(state: AgentState, config: RunnableConfig) -> AgentState:
+    llama_guard = LlamaGuard()
+    safety_output = await llama_guard.ainvoke("User", state["messages"])
+    return {"safety": safety_output, "messages": []}
+
+
+async def block_unsafe_content(state: AgentState, config: RunnableConfig) -> AgentState:
+    safety: LlamaGuardOutput = state["safety"]
+    return {"messages": [format_safety_message(safety)]}
+
+
+# Define the graph
+agent = StateGraph(AgentState)
+agent.add_node("model", acall_model)
+agent.add_node("tools", ToolNode(tools))
+agent.add_node("guard_input", llama_guard_input)
+agent.add_node("block_unsafe_content", block_unsafe_content)
+agent.set_entry_point("guard_input")
+
+
+# Check for unsafe input and block further processing if found
+def check_safety(state: AgentState) -> Literal["unsafe", "safe"]:
+    safety: LlamaGuardOutput = state["safety"]
+    match safety.safety_assessment:
+        case SafetyAssessment.UNSAFE:
+            return "unsafe"
+        case _:
+            return "safe"
+
+
+agent.add_conditional_edges(
+    "guard_input", check_safety, {"unsafe": "block_unsafe_content", "safe": "model"}
+)
+
+# Always END after blocking unsafe content
+agent.add_edge("block_unsafe_content", END)
+
+# Always run "model" after "tools"
+agent.add_edge("tools", "model")
+
+
+# After "model", if there are tool calls, run "tools". Otherwise END.
+def pending_tool_calls(state: AgentState) -> Literal["tools", "done"]:
+    last_message = state["messages"][-1]
+    if not isinstance(last_message, AIMessage):
+        raise TypeError(f"Expected AIMessage, got {type(last_message)}")
+    if last_message.tool_calls:
+        return "tools"
+    return "done"
+
+
+agent.add_conditional_edges("model", pending_tool_calls, {"tools": "tools", "done": END})
+
+rag_assistant = agent.compile()

src/agents/research_assistant.py

@@ -0,0 +1,148 @@
+from datetime import datetime
+from typing import Literal
+
+from langchain_community.tools import DuckDuckGoSearchResults, OpenWeatherMapQueryRun
+from langchain_community.utilities import OpenWeatherMapAPIWrapper
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, SystemMessage
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnableSerializable
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.managed import RemainingSteps
+from langgraph.prebuilt import ToolNode
+
+from agents.llama_guard import LlamaGuard, LlamaGuardOutput, SafetyAssessment
+from agents.tools import calculator
+from core import get_model, settings
+
+
+class AgentState(MessagesState, total=False):
+    """`total=False` is PEP589 specs.
+
+    documentation: https://typing.readthedocs.io/en/latest/spec/typeddict.html#totality
+    """
+
+    safety: LlamaGuardOutput
+    remaining_steps: RemainingSteps
+
+
+web_search = DuckDuckGoSearchResults(name="WebSearch")
+tools = [web_search, calculator]
+
+# Add weather tool if API key is set
+# Register for an API key at https://openweathermap.org/api/
+if settings.OPENWEATHERMAP_API_KEY:
+    wrapper = OpenWeatherMapAPIWrapper(
+        openweathermap_api_key=settings.OPENWEATHERMAP_API_KEY.get_secret_value()
+    )
+    tools.append(OpenWeatherMapQueryRun(name="Weather", api_wrapper=wrapper))
+
+current_date = datetime.now().strftime("%B %d, %Y")
+instructions = f"""
+    You are a helpful research assistant with the ability to search the web and use other tools.
+    Today's date is {current_date}.
+
+    NOTE: THE USER CAN'T SEE THE TOOL RESPONSE.
+
+    A few things to remember:
+    - Please include markdown-formatted links to any citations used in your response. Only include one
+    or two citations per response unless more are needed. ONLY USE LINKS RETURNED BY THE TOOLS.
+    - Use calculator tool with numexpr to answer math questions. The user does not understand numexpr,
+      so for the final response, use human readable format - e.g. "300 * 200", not "(300 \\times 200)".
+    """
+
+
+def wrap_model(model: BaseChatModel) -> RunnableSerializable[AgentState, AIMessage]:
+    bound_model = model.bind_tools(tools)
+    preprocessor = RunnableLambda(
+        lambda state: [SystemMessage(content=instructions)] + state["messages"],
+        name="StateModifier",
+    )
+    return preprocessor | bound_model  # type: ignore[return-value]
+
+
+def format_safety_message(safety: LlamaGuardOutput) -> AIMessage:
+    content = (
+        f"This conversation was flagged for unsafe content: {', '.join(safety.unsafe_categories)}"
+    )
+    return AIMessage(content=content)
+
+
+async def acall_model(state: AgentState, config: RunnableConfig) -> AgentState:
+    m = get_model(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(m)
+    response = await model_runnable.ainvoke(state, config)
+
+    # Run llama guard check here to avoid returning the message if it's unsafe
+    llama_guard = LlamaGuard()
+    safety_output = await llama_guard.ainvoke("Agent", state["messages"] + [response])
+    if safety_output.safety_assessment == SafetyAssessment.UNSAFE:
+        return {"messages": [format_safety_message(safety_output)], "safety": safety_output}
+
+    if state["remaining_steps"] < 2 and response.tool_calls:
+        return {
+            "messages": [
+                AIMessage(
+                    id=response.id,
+                    content="Sorry, need more steps to process this request.",
+                )
+            ]
+        }
+    # We return a list, because this will get added to the existing list
+    return {"messages": [response]}
+
+
+async def llama_guard_input(state: AgentState, config: RunnableConfig) -> AgentState:
+    llama_guard = LlamaGuard()
+    safety_output = await llama_guard.ainvoke("User", state["messages"])
+    return {"safety": safety_output, "messages": []}
+
+
+async def block_unsafe_content(state: AgentState, config: RunnableConfig) -> AgentState:
+    safety: LlamaGuardOutput = state["safety"]
+    return {"messages": [format_safety_message(safety)]}
+
+
+# Define the graph
+agent = StateGraph(AgentState)
+agent.add_node("model", acall_model)
+agent.add_node("tools", ToolNode(tools))
+agent.add_node("guard_input", llama_guard_input)
+agent.add_node("block_unsafe_content", block_unsafe_content)
+agent.set_entry_point("guard_input")
+
+
+# Check for unsafe input and block further processing if found
+def check_safety(state: AgentState) -> Literal["unsafe", "safe"]:
+    safety: LlamaGuardOutput = state["safety"]
+    match safety.safety_assessment:
+        case SafetyAssessment.UNSAFE:
+            return "unsafe"
+        case _:
+            return "safe"
+
+
+agent.add_conditional_edges(
+    "guard_input", check_safety, {"unsafe": "block_unsafe_content", "safe": "model"}
+)
+
+# Always END after blocking unsafe content
+agent.add_edge("block_unsafe_content", END)
+
+# Always run "model" after "tools"
+agent.add_edge("tools", "model")
+
+
+# After "model", if there are tool calls, run "tools". Otherwise END.
+def pending_tool_calls(state: AgentState) -> Literal["tools", "done"]:
+    last_message = state["messages"][-1]
+    if not isinstance(last_message, AIMessage):
+        raise TypeError(f"Expected AIMessage, got {type(last_message)}")
+    if last_message.tool_calls:
+        return "tools"
+    return "done"
+
+
+agent.add_conditional_edges("model", pending_tool_calls, {"tools": "tools", "done": END})
+
+
+research_assistant = agent.compile()

src/agents/tools.py

@@ -0,0 +1,56 @@
+import re
+import math
+import numexpr
+
+from memory.postgres import load_pgvector_retriever
+from langchain_core.tools import BaseTool, tool
+
+
+def calculator_func(expression: str) -> str:
+    """Calculates a math expression using numexpr.
+
+    Useful for when you need to answer questions about math using numexpr.
+    This tool is only for math questions and nothing else. Only input
+    math expressions.
+
+    Args:
+        expression (str): A valid numexpr formatted math expression.
+
+    Returns:
+        str: The result of the math expression.
+    """
+
+    try:
+        local_dict = {"pi": math.pi, "e": math.e}
+        output = str(
+            numexpr.evaluate(
+                expression.strip(),
+                global_dict={},  # restrict access to globals
+                local_dict=local_dict,  # add common mathematical functions
+            )
+        )
+        return re.sub(r"^\[|\]$", "", output)
+    except Exception as e:
+        raise ValueError(
+            f'calculator("{expression}") raised error: {e}.'
+            " Please try again with a valid numerical expression"
+        )
+
+
+calculator: BaseTool = tool(calculator_func)
+calculator.name = "Calculator"
+
+def format_contexts(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+def database_search_func(query: str) -> str:
+    """Searches the vector DB for information in the portfolio."""
+    retriever = load_pgvector_retriever()
+    documents = retriever.invoke(query)
+    context_str = format_contexts(documents)
+
+    return context_str
+
+
+database_search: BaseTool = tool(database_search_func)
+database_search.name = "Database_Search"

src/agents/utils.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from langchain_core.messages import ChatMessage
+from langgraph.types import StreamWriter
+from pydantic import BaseModel, Field
+
+
+class CustomData(BaseModel):
+    "Custom data being sent by an agent"
+
+    data: dict[str, Any] = Field(description="The custom data")
+
+    def to_langchain(self) -> ChatMessage:
+        return ChatMessage(content=[self.data], role="custom")
+
+    def dispatch(self, writer: StreamWriter) -> None:
+        writer(self.to_langchain())

src/core/__init__.py

@@ -0,0 +1,4 @@
+from core.llm import get_model
+from core.settings import settings
+
+__all__ = ["settings", "get_model"]

src/core/embeddings.py

@@ -0,0 +1,37 @@
+from functools import cache
+from typing import TypeAlias
+
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_ollama import OllamaEmbeddings
+from langchain_openai import OpenAIEmbeddings
+
+from core.settings import settings
+from schema.models import (
+    AllEmbeddingModelEnum,
+    GoogleEmbeddingModelName,
+    OllamaEmbeddingModelName,
+    OpenAIEmbeddingModelName,
+)
+
+EmbeddingT: TypeAlias = (
+    OpenAIEmbeddings
+    | GoogleGenerativeAIEmbeddings
+    | OllamaEmbeddings
+)
+
+
+@cache
+def get_embeddings(model_name: AllEmbeddingModelEnum, /) -> EmbeddingT:
+    if model_name in OpenAIEmbeddingModelName:
+        return OpenAIEmbeddings(model=model_name.value)
+    
+    if model_name in GoogleEmbeddingModelName:
+        return GoogleGenerativeAIEmbeddings(model=model_name.value)
+    
+    if model_name in OllamaEmbeddingModelName:
+        return OllamaEmbeddings(
+            model=settings.OLLAMA_EMBEDDING_MODEL or model_name.value,
+            base_url=settings.OLLAMA_BASE_URL,
+        )
+    
+    raise ValueError(f"Unsupported embedding model: {model_name}")

src/core/llm.py

@@ -0,0 +1,147 @@
+from functools import cache
+from typing import TypeAlias
+
+from langchain_anthropic import ChatAnthropic
+from langchain_aws import ChatBedrock
+from langchain_community.chat_models import FakeListChatModel
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_google_vertexai import ChatVertexAI
+from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
+from langchain_openai import AzureChatOpenAI, ChatOpenAI
+
+from core.settings import settings
+from schema.models import (
+    AllModelEnum,
+    AnthropicModelName,
+    AWSModelName,
+    AzureOpenAIModelName,
+    DeepseekModelName,
+    FakeModelName,
+    GoogleModelName,
+    GroqModelName,
+    OllamaModelName,
+    OpenAICompatibleName,
+    OpenAIModelName,
+    OpenRouterModelName,
+    VertexAIModelName,
+)
+
+_MODEL_TABLE = (
+    {m: m.value for m in OpenAIModelName}
+    | {m: m.value for m in OpenAICompatibleName}
+    | {m: m.value for m in AzureOpenAIModelName}
+    | {m: m.value for m in DeepseekModelName}
+    | {m: m.value for m in AnthropicModelName}
+    | {m: m.value for m in GoogleModelName}
+    | {m: m.value for m in VertexAIModelName}
+    | {m: m.value for m in GroqModelName}
+    | {m: m.value for m in AWSModelName}
+    | {m: m.value for m in OllamaModelName}
+    | {m: m.value for m in OpenRouterModelName}
+    | {m: m.value for m in FakeModelName}
+)
+
+
+class FakeToolModel(FakeListChatModel):
+    def __init__(self, responses: list[str]):
+        super().__init__(responses=responses)
+
+    def bind_tools(self, tools):
+        return self
+
+
+ModelT: TypeAlias = (
+    AzureChatOpenAI
+    | ChatOpenAI
+    | ChatAnthropic
+    | ChatGoogleGenerativeAI
+    | ChatVertexAI
+    | ChatGroq
+    | ChatBedrock
+    | ChatOllama
+    | FakeToolModel
+)
+
+
+@cache
+def get_model(model_name: AllModelEnum, /) -> ModelT:
+    # NOTE: models with streaming=True will send tokens as they are generated
+    # if the /stream endpoint is called with stream_tokens=True (the default)
+    api_model_name = _MODEL_TABLE.get(model_name)
+    if not api_model_name:
+        raise ValueError(f"Unsupported model: {model_name}")
+
+    if model_name in OpenAIModelName:
+        return ChatOpenAI(model=api_model_name, streaming=True)
+    if model_name in OpenAICompatibleName:
+        if not settings.COMPATIBLE_BASE_URL or not settings.COMPATIBLE_MODEL:
+            raise ValueError("OpenAICompatible base url and endpoint must be configured")
+
+        return ChatOpenAI(
+            model=settings.COMPATIBLE_MODEL,
+            temperature=0.5,
+            streaming=True,
+            openai_api_base=settings.COMPATIBLE_BASE_URL,
+            openai_api_key=settings.COMPATIBLE_API_KEY,
+        )
+    if model_name in AzureOpenAIModelName:
+        if not settings.AZURE_OPENAI_API_KEY or not settings.AZURE_OPENAI_ENDPOINT:
+            raise ValueError("Azure OpenAI API key and endpoint must be configured")
+
+        return AzureChatOpenAI(
+            azure_endpoint=settings.AZURE_OPENAI_ENDPOINT,
+            deployment_name=api_model_name,
+            api_version=settings.AZURE_OPENAI_API_VERSION,
+            temperature=0.5,
+            streaming=True,
+            timeout=60,
+            max_retries=3,
+        )
+    if model_name in DeepseekModelName:
+        return ChatOpenAI(
+            model=api_model_name,
+            temperature=0.5,
+            streaming=True,
+            openai_api_base="https://api.deepseek.com",
+            openai_api_key=settings.DEEPSEEK_API_KEY,
+        )
+    if model_name in AnthropicModelName:
+        return ChatAnthropic(model=api_model_name, temperature=0.5, streaming=True)
+    if model_name in GoogleModelName:
+        return ChatGoogleGenerativeAI(model=api_model_name, temperature=0.5, streaming=True)
+    if model_name in VertexAIModelName:
+        return ChatVertexAI(model=api_model_name, temperature=0.5, streaming=True)
+    if model_name in GroqModelName:
+        # Guard and safeguard models should use temperature=0.0 for deterministic outputs
+        guard_models = {
+            GroqModelName.LLAMA_GUARD_4_12B,
+            GroqModelName.LLAMA_PROMPT_GUARD_2_22M,
+            GroqModelName.LLAMA_PROMPT_GUARD_2_86M,
+            GroqModelName.OPENAI_GPT_OSS_SAFEGUARD_20B,
+        }
+        if model_name in guard_models:
+            return ChatGroq(model=api_model_name, temperature=0.0)  # type: ignore[call-arg]
+        return ChatGroq(model=api_model_name, temperature=0.5)  # type: ignore[call-arg]
+    if model_name in AWSModelName:
+        return ChatBedrock(model_id=api_model_name, temperature=0.5)
+    if model_name in OllamaModelName:
+        if settings.OLLAMA_BASE_URL:
+            chat_ollama = ChatOllama(
+                model=settings.OLLAMA_MODEL, temperature=0.5, base_url=settings.OLLAMA_BASE_URL
+            )
+        else:
+            chat_ollama = ChatOllama(model=settings.OLLAMA_MODEL, temperature=0.5)
+        return chat_ollama
+    if model_name in OpenRouterModelName:
+        return ChatOpenAI(
+            model=api_model_name,
+            temperature=0.5,
+            streaming=True,
+            base_url="https://openrouter.ai/api/v1/",
+            api_key=settings.OPENROUTER_API_KEY,
+        )
+    if model_name in FakeModelName:
+        return FakeToolModel(responses=["This is a test response from the fake model."])
+
+    raise ValueError(f"Unsupported model: {model_name}")

src/core/settings.py

@@ -0,0 +1,289 @@
+from enum import StrEnum
+from json import loads
+from typing import Annotated, Any
+
+from dotenv import find_dotenv
+from pydantic import (
+    BeforeValidator,
+    Field,
+    HttpUrl,
+    SecretStr,
+    TypeAdapter,
+    computed_field,
+)
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from schema.models import (
+    AllModelEnum,
+    AnthropicModelName,
+    AWSModelName,
+    AzureOpenAIModelName,
+    DeepseekModelName,
+    FakeModelName,
+    GoogleModelName,
+    GroqModelName,
+    OllamaModelName,
+    OpenAICompatibleName,
+    OpenAIModelName,
+    OpenRouterModelName,
+    Provider,
+    VertexAIModelName,
+    AllEmbeddingModelEnum,
+    OpenAIEmbeddingModelName,
+    GoogleEmbeddingModelName,
+    OllamaEmbeddingModelName,
+)
+
+
+class DatabaseType(StrEnum):
+    SQLITE = "sqlite"
+    POSTGRES = "postgres"
+    MONGO = "mongo"
+
+
+class LogLevel(StrEnum):
+    DEBUG = "DEBUG"
+    INFO = "INFO"
+    WARNING = "WARNING"
+    ERROR = "ERROR"
+    CRITICAL = "CRITICAL"
+
+    def to_logging_level(self) -> int:
+        """Convert to Python logging level constant."""
+        import logging
+
+        mapping = {
+            LogLevel.DEBUG: logging.DEBUG,
+            LogLevel.INFO: logging.INFO,
+            LogLevel.WARNING: logging.WARNING,
+            LogLevel.ERROR: logging.ERROR,
+            LogLevel.CRITICAL: logging.CRITICAL,
+        }
+        return mapping[self]
+
+
+def check_str_is_http(x: str) -> str:
+    http_url_adapter = TypeAdapter(HttpUrl)
+    return str(http_url_adapter.validate_python(x))
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_file=find_dotenv(),
+        env_file_encoding="utf-8",
+        env_ignore_empty=True,
+        extra="ignore",
+        validate_default=False,
+    )
+    MODE: str | None = None
+
+    HOST: str = "0.0.0.0"
+    PORT: int = 7860
+    GRACEFUL_SHUTDOWN_TIMEOUT: int = 30
+    LOG_LEVEL: LogLevel = LogLevel.WARNING
+
+    AUTH_SECRET: SecretStr | None = None
+    CORS_ORIGINS: Annotated[Any, BeforeValidator(lambda x: x.split(",") if isinstance(x, str) else x)] = [
+        "http://localhost:3000",
+        "http://localhost:8081",
+        "http://localhost:5173",
+    ]
+
+    OPENAI_API_KEY: SecretStr | None = None
+    DEEPSEEK_API_KEY: SecretStr | None = None
+    ANTHROPIC_API_KEY: SecretStr | None = None
+    GOOGLE_API_KEY: SecretStr | None = None
+    GOOGLE_APPLICATION_CREDENTIALS: SecretStr | None = None
+    GROQ_API_KEY: SecretStr | None = None
+    USE_AWS_BEDROCK: bool = False
+    OLLAMA_MODEL: str | None = None
+    OLLAMA_BASE_URL: str | None = None
+    USE_FAKE_MODEL: bool = False
+    OPENROUTER_API_KEY: str | None = None
+
+    # If DEFAULT_MODEL is None, it will be set in model_post_init
+    DEFAULT_MODEL: AllModelEnum | None = None  # type: ignore[assignment]
+    AVAILABLE_MODELS: set[AllModelEnum] = set()  # type: ignore[assignment]
+
+    # Embedding Settings
+    DEFAULT_EMBEDDING_MODEL: AllEmbeddingModelEnum | None = None  # type: ignore[assignment]
+    AVAILABLE_EMBEDDING_MODELS: set[AllEmbeddingModelEnum] = set()  # type: ignore[assignment]
+    OLLAMA_EMBEDDING_MODEL: str | None = None
+
+    # Set openai compatible api, mainly used for proof of concept
+    COMPATIBLE_MODEL: str | None = None
+    COMPATIBLE_API_KEY: SecretStr | None = None
+    COMPATIBLE_BASE_URL: str | None = None
+
+    OPENWEATHERMAP_API_KEY: SecretStr | None = None
+
+    # MCP Configuration
+    GITHUB_PAT: SecretStr | None = None
+    MCP_GITHUB_SERVER_URL: str = "https://api.githubcopilot.com/mcp/"
+
+    LANGCHAIN_TRACING_V2: bool = False
+    LANGCHAIN_PROJECT: str = "default"
+    LANGCHAIN_ENDPOINT: Annotated[str, BeforeValidator(check_str_is_http)] = (
+        "https://api.smith.langchain.com"
+    )
+    LANGCHAIN_API_KEY: SecretStr | None = None
+
+    LANGFUSE_TRACING: bool = False
+    LANGFUSE_HOST: Annotated[str, BeforeValidator(check_str_is_http)] = "https://cloud.langfuse.com"
+    LANGFUSE_PUBLIC_KEY: SecretStr | None = None
+    LANGFUSE_SECRET_KEY: SecretStr | None = None
+
+    # Database Configuration
+    DATABASE_TYPE: DatabaseType = (
+        DatabaseType.SQLITE
+    )  # Options: DatabaseType.SQLITE or DatabaseType.POSTGRES
+    SQLITE_DB_PATH: str = "checkpoints.db"
+
+    # PostgreSQL Configuration
+    POSTGRES_USER: str | None = None
+    POSTGRES_PASSWORD: SecretStr | None = None
+    POSTGRES_HOST: str | None = None
+    POSTGRES_PORT: int | None = None
+    POSTGRES_DB: str | None = None
+    POSTGRES_APPLICATION_NAME: str = "agent-service-toolkit"
+    POSTGRES_MIN_CONNECTIONS_PER_POOL: int = 1
+    POSTGRES_MAX_CONNECTIONS_PER_POOL: int = 1
+    VECTOR_STORE_COLLECTION_NAME: str = "vector_store"
+
+    # MongoDB Configuration
+    MONGO_HOST: str | None = None
+    MONGO_PORT: int | None = None
+    MONGO_DB: str | None = None
+    MONGO_USER: str | None = None
+    MONGO_PASSWORD: SecretStr | None = None
+    MONGO_AUTH_SOURCE: str | None = None
+
+    # Azure OpenAI Settings
+    AZURE_OPENAI_API_KEY: SecretStr | None = None
+    AZURE_OPENAI_ENDPOINT: str | None = None
+    AZURE_OPENAI_API_VERSION: str = "2024-02-15-preview"
+    AZURE_OPENAI_DEPLOYMENT_MAP: dict[str, str] = Field(
+        default_factory=dict, description="Map of model names to Azure deployment IDs"
+    )
+
+    def model_post_init(self, __context: Any) -> None:
+        api_keys = {
+            Provider.OPENAI: self.OPENAI_API_KEY,
+            Provider.OPENAI_COMPATIBLE: self.COMPATIBLE_BASE_URL and self.COMPATIBLE_MODEL,
+            Provider.DEEPSEEK: self.DEEPSEEK_API_KEY,
+            Provider.ANTHROPIC: self.ANTHROPIC_API_KEY,
+            Provider.GOOGLE: self.GOOGLE_API_KEY,
+            Provider.VERTEXAI: self.GOOGLE_APPLICATION_CREDENTIALS,
+            Provider.GROQ: self.GROQ_API_KEY,
+            Provider.AWS: self.USE_AWS_BEDROCK,
+            Provider.OLLAMA: self.OLLAMA_MODEL,
+            Provider.FAKE: self.USE_FAKE_MODEL,
+            Provider.AZURE_OPENAI: self.AZURE_OPENAI_API_KEY,
+            Provider.OPENROUTER: self.OPENROUTER_API_KEY,
+        }
+        active_keys = [k for k, v in api_keys.items() if v]
+        if not active_keys:
+            raise ValueError("At least one LLM API key must be provided.")
+
+        for provider in active_keys:
+            match provider:
+                case Provider.OPENAI:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = OpenAIModelName.GPT_5_NANO
+                    self.AVAILABLE_MODELS.update(set(OpenAIModelName))
+                case Provider.OPENAI_COMPATIBLE:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = OpenAICompatibleName.OPENAI_COMPATIBLE
+                    self.AVAILABLE_MODELS.update(set(OpenAICompatibleName))
+                case Provider.DEEPSEEK:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = DeepseekModelName.DEEPSEEK_CHAT
+                    self.AVAILABLE_MODELS.update(set(DeepseekModelName))
+                case Provider.ANTHROPIC:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = AnthropicModelName.HAIKU_45
+                    self.AVAILABLE_MODELS.update(set(AnthropicModelName))
+                case Provider.GOOGLE:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = GoogleModelName.GEMINI_20_FLASH
+                    self.AVAILABLE_MODELS.update(set(GoogleModelName))
+                case Provider.VERTEXAI:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = VertexAIModelName.GEMINI_20_FLASH
+                    self.AVAILABLE_MODELS.update(set(VertexAIModelName))
+                case Provider.GROQ:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = GroqModelName.LLAMA_31_8B_INSTANT
+                    self.AVAILABLE_MODELS.update(set(GroqModelName))
+                case Provider.AWS:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = AWSModelName.BEDROCK_HAIKU
+                    self.AVAILABLE_MODELS.update(set(AWSModelName))
+                case Provider.OLLAMA:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = OllamaModelName.OLLAMA_GENERIC
+                    self.AVAILABLE_MODELS.update(set(OllamaModelName))
+                case Provider.OPENROUTER:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = OpenRouterModelName.GEMINI_25_FLASH
+                    self.AVAILABLE_MODELS.update(set(OpenRouterModelName))
+                case Provider.FAKE:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = FakeModelName.FAKE
+                    self.AVAILABLE_MODELS.update(set(FakeModelName))
+                case Provider.AZURE_OPENAI:
+                    if self.DEFAULT_MODEL is None:
+                        self.DEFAULT_MODEL = AzureOpenAIModelName.AZURE_GPT_4O_MINI
+                    self.AVAILABLE_MODELS.update(set(AzureOpenAIModelName))
+                    # Validate Azure OpenAI settings if Azure provider is available
+                    if not self.AZURE_OPENAI_API_KEY:
+                        raise ValueError("AZURE_OPENAI_API_KEY must be set")
+                    if not self.AZURE_OPENAI_ENDPOINT:
+                        raise ValueError("AZURE_OPENAI_ENDPOINT must be set")
+                    if not self.AZURE_OPENAI_DEPLOYMENT_MAP:
+                        raise ValueError("AZURE_OPENAI_DEPLOYMENT_MAP must be set")
+
+                    # Parse deployment map if it's a string
+                    if isinstance(self.AZURE_OPENAI_DEPLOYMENT_MAP, str):
+                        try:
+                            self.AZURE_OPENAI_DEPLOYMENT_MAP = loads(
+                                self.AZURE_OPENAI_DEPLOYMENT_MAP
+                            )
+                        except Exception as e:
+                            raise ValueError(f"Invalid AZURE_OPENAI_DEPLOYMENT_MAP JSON: {e}")
+
+                    # Validate required deployments exist
+                    required_models = {"gpt-4o", "gpt-4o-mini"}
+                    missing_models = required_models - set(self.AZURE_OPENAI_DEPLOYMENT_MAP.keys())
+                    if missing_models:
+                        raise ValueError(f"Missing required Azure deployments: {missing_models}")
+                case _:
+                    raise ValueError(f"Unknown provider: {provider}")
+
+        for provider in active_keys:
+            match provider:
+                case Provider.OPENAI:
+                    if self.DEFAULT_EMBEDDING_MODEL is None:
+                        self.DEFAULT_EMBEDDING_MODEL = OpenAIEmbeddingModelName.TEXT_EMBEDDING_3_SMALL
+                    self.AVAILABLE_EMBEDDING_MODELS.update(set(OpenAIEmbeddingModelName))
+                case Provider.GOOGLE:
+                    if self.DEFAULT_EMBEDDING_MODEL is None:
+                        self.DEFAULT_EMBEDDING_MODEL = GoogleEmbeddingModelName.TEXT_EMBEDDING_004
+                    self.AVAILABLE_EMBEDDING_MODELS.update(set(GoogleEmbeddingModelName))
+                case Provider.OLLAMA:
+                    if self.DEFAULT_EMBEDDING_MODEL is None:
+                        self.DEFAULT_EMBEDDING_MODEL = OllamaEmbeddingModelName.NOMIC_EMBED_TEXT
+                    self.AVAILABLE_EMBEDDING_MODELS.update(set(OllamaEmbeddingModelName))
+                    if not self.OLLAMA_EMBEDDING_MODEL:
+                        self.OLLAMA_EMBEDDING_MODEL = OllamaEmbeddingModelName.NOMIC_EMBED_TEXT
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def BASE_URL(self) -> str:
+        return f"http://{self.HOST}:{self.PORT}"
+
+    def is_dev(self) -> bool:
+        return self.MODE == "dev"
+
+
+settings = Settings()

src/memory/__init__.py

@@ -0,0 +1,40 @@
+from contextlib import AbstractAsyncContextManager
+
+from langgraph.checkpoint.mongodb.aio import AsyncMongoDBSaver
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+
+from core.settings import DatabaseType, settings
+from memory.mongodb import get_mongo_saver
+from memory.postgres import get_postgres_saver, get_postgres_store
+from memory.sqlite import get_sqlite_saver, get_sqlite_store
+
+
+def initialize_database() -> AbstractAsyncContextManager[
+    AsyncSqliteSaver | AsyncPostgresSaver | AsyncMongoDBSaver
+]:
+    """
+    Initialize the appropriate database checkpointer based on configuration.
+    Returns an initialized AsyncCheckpointer instance.
+    """
+    if settings.DATABASE_TYPE == DatabaseType.POSTGRES:
+        return get_postgres_saver()
+    if settings.DATABASE_TYPE == DatabaseType.MONGO:
+        return get_mongo_saver()
+    else:  # Default to SQLite
+        return get_sqlite_saver()
+
+
+def initialize_store():
+    """
+    Initialize the appropriate store based on configuration.
+    Returns an async context manager for the initialized store.
+    """
+    if settings.DATABASE_TYPE == DatabaseType.POSTGRES:
+        return get_postgres_store()
+    # TODO: Add Mongo store - https://pypi.org/project/langgraph-store-mongodb/
+    else:  # Default to SQLite
+        return get_sqlite_store()
+
+
+__all__ = ["initialize_database", "initialize_store"]

src/memory/mongodb.py

@@ -0,0 +1,62 @@
+import logging
+import urllib.parse
+from contextlib import AbstractAsyncContextManager
+
+from langgraph.checkpoint.mongodb.aio import AsyncMongoDBSaver
+
+from core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+
+def _has_auth_credentials() -> bool:
+    required_auth = ["MONGO_USER", "MONGO_PASSWORD", "MONGO_AUTH_SOURCE"]
+    set_auth = [var for var in required_auth if getattr(settings, var, None)]
+    if len(set_auth) > 0 and len(set_auth) != len(required_auth):
+        raise ValueError(
+            f"If any of the following environment variables are set, all must be set: {', '.join(required_auth)}."
+        )
+    return len(set_auth) == len(required_auth)
+
+
+def validate_mongo_config() -> None:
+    """
+    Validate that all required MongoDB configuration is present.
+    Raises ValueError if any required configuration is missing.
+    """
+    required_always = ["MONGO_HOST", "MONGO_PORT", "MONGO_DB"]
+    missing_always = [var for var in required_always if not getattr(settings, var, None)]
+    if missing_always:
+        raise ValueError(
+            f"Missing required MongoDB configuration: {', '.join(missing_always)}. "
+            "These environment variables must be set to use MongoDB persistence."
+        )
+
+    _has_auth_credentials()
+
+
+def get_mongo_connection_string() -> str:
+    """Build and return the MongoDB connection string from settings."""
+
+    if _has_auth_credentials():
+        if settings.MONGO_PASSWORD is None:  # for type checking
+            raise ValueError("MONGO_PASSWORD is not set")
+        password = settings.MONGO_PASSWORD.get_secret_value().strip()
+        password_escaped = urllib.parse.quote_plus(password)
+        return (
+            f"mongodb://{settings.MONGO_USER}:{password_escaped}@"
+            f"{settings.MONGO_HOST}:{settings.MONGO_PORT}/"
+            f"?authSource={settings.MONGO_AUTH_SOURCE}"
+        )
+    else:
+        return f"mongodb://{settings.MONGO_HOST}:{settings.MONGO_PORT}/"
+
+
+def get_mongo_saver() -> AbstractAsyncContextManager[AsyncMongoDBSaver]:
+    """Initialize and return a MongoDB saver instance."""
+    validate_mongo_config()
+    if settings.MONGO_DB is None:  # for type checking
+        raise ValueError("MONGO_DB is not set")
+    return AsyncMongoDBSaver.from_conn_string(
+        get_mongo_connection_string(), db_name=settings.MONGO_DB
+    )

src/memory/postgres.py

@@ -0,0 +1,135 @@
+import logging
+from contextlib import asynccontextmanager
+
+from core.embeddings import get_embeddings
+from langchain_postgres import PGVector
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.store.postgres import AsyncPostgresStore
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
+
+from core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+
+def validate_postgres_config() -> None:
+    """
+    Validate that all required PostgreSQL configuration is present.
+    Raises ValueError if any required configuration is missing.
+    """
+
+    required_vars = [
+        "POSTGRES_USER",
+        "POSTGRES_PASSWORD",
+        "POSTGRES_HOST",
+        "POSTGRES_PORT",
+        "POSTGRES_DB",
+    ]
+
+    missing = [var for var in required_vars if not getattr(settings, var, None)]
+    if missing:
+        raise ValueError(
+            f"Missing required PostgreSQL configuration: {', '.join(missing)}. "
+            "All individual POSTGRES_* environment variables must be set to use PostgreSQL persistence."
+        )
+
+    if settings.POSTGRES_MIN_CONNECTIONS_PER_POOL > settings.POSTGRES_MAX_CONNECTIONS_PER_POOL:
+        raise ValueError(
+            f"POSTGRES_MIN_CONNECTIONS_PER_POOL ({settings.POSTGRES_MIN_CONNECTIONS_PER_POOL}) must be less than or equal to POSTGRES_MAX_CONNECTIONS_PER_POOL ({settings.POSTGRES_MAX_CONNECTIONS_PER_POOL})"
+        )
+
+
+def get_postgres_connection_string() -> str:
+    """Build and return the PostgreSQL connection string from settings."""
+    if settings.POSTGRES_PASSWORD is None:
+        raise ValueError("POSTGRES_PASSWORD is not set")
+    return (
+        f"postgresql://{settings.POSTGRES_USER}:"
+        f"{settings.POSTGRES_PASSWORD.get_secret_value()}@"
+        f"{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/"
+        f"{settings.POSTGRES_DB}/?sslmode=require"
+    )
+
+
+@asynccontextmanager
+async def get_postgres_saver():
+    """Initialize and return a PostgreSQL saver instance based on a connection pool for more resilient connections."""
+    validate_postgres_config()
+    application_name = settings.POSTGRES_APPLICATION_NAME + "-" + "saver"
+
+    async with AsyncConnectionPool(
+        get_postgres_connection_string(),
+        min_size=settings.POSTGRES_MIN_CONNECTIONS_PER_POOL,
+        max_size=settings.POSTGRES_MAX_CONNECTIONS_PER_POOL,
+        # Langgraph requires autocommmit=true and row_factory to be set to dict_row.
+        # Application_name is passed so you can identify the connection in your Postgres database connection manager.
+        kwargs={"autocommit": True, "row_factory": dict_row, "application_name": application_name},
+        # makes sure that the connection is still valid before using it
+        check=AsyncConnectionPool.check_connection,
+    ) as pool:
+        try:
+            checkpointer = AsyncPostgresSaver(pool)
+            await checkpointer.setup()
+            yield checkpointer
+        finally:
+            await pool.close()
+
+
+@asynccontextmanager
+async def get_postgres_store():
+    """
+    Get a PostgreSQL store instance based on a connection pool for more resilient connections.
+
+    Returns an AsyncPostgresStore instance that can be used with async context manager pattern.
+
+    """
+    validate_postgres_config()
+    application_name = settings.POSTGRES_APPLICATION_NAME + "-" + "store"
+
+    async with AsyncConnectionPool(
+        get_postgres_connection_string(),
+        min_size=settings.POSTGRES_MIN_CONNECTIONS_PER_POOL,
+        max_size=settings.POSTGRES_MAX_CONNECTIONS_PER_POOL,
+        # Langgraph requires autocommmit=true and row_factory to be set to dict_row
+        # Application_name is passed so you can identify the connection in your Postgres database connection manager.
+        kwargs={"autocommit": True, "row_factory": dict_row, "application_name": application_name},
+        # makes sure that the connection is still valid before using it
+        check=AsyncConnectionPool.check_connection,
+    ) as pool:
+        try:
+            store = AsyncPostgresStore(pool)
+            await store.setup()
+            yield store
+        finally:
+            await pool.close()
+
+def get_pgvector_connection_string() -> str:
+    """Build and return the PostgreSQL connection string for vectors from settings."""
+    return (
+        f"postgresql+psycopg://{settings.POSTGRES_USER}:"
+        f"{settings.POSTGRES_PASSWORD.get_secret_value()}@"
+        f"{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/"
+        f"{settings.POSTGRES_DB}?sslmode=require"
+    )
+
+def load_pgvector_store():
+    """Get a PostgreSQL vectors store instance."""
+    validate_postgres_config()
+
+    return PGVector(
+        connection=get_pgvector_connection_string(),
+        collection_name=settings.VECTOR_STORE_COLLECTION_NAME,
+        embeddings=get_embeddings(settings.DEFAULT_EMBEDDING_MODEL),
+    )
+
+def load_pgvector_retriever(k: int = 6):
+    store = load_pgvector_store()
+    return store.as_retriever(
+        search_type="mmr",
+        search_kwargs={
+            "k": k,
+            "fetch_k": 20,                # candidates
+            "lambda_mult": 0.6,           # relevance vs diversity
+        },
+    )

src/memory/sqlite.py

@@ -0,0 +1,40 @@
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+from langgraph.store.memory import InMemoryStore
+
+from core.settings import settings
+
+
+def get_sqlite_saver() -> AbstractAsyncContextManager[AsyncSqliteSaver]:
+    """Initialize and return a SQLite saver instance."""
+    return AsyncSqliteSaver.from_conn_string(settings.SQLITE_DB_PATH)
+
+
+class AsyncInMemoryStore:
+    """Wrapper for InMemoryStore that provides an async context manager interface."""
+
+    def __init__(self):
+        self.store = InMemoryStore()
+
+    async def __aenter__(self):
+        return self.store
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        # No cleanup needed for InMemoryStore
+        pass
+
+    async def setup(self):
+        # No-op method for compatibility with PostgresStore
+        pass
+
+
+@asynccontextmanager
+async def get_sqlite_store():
+    """Initialize and return a store instance for long-term memory.
+
+    Note: SQLite-specific store isn't available in LangGraph,
+    so we use InMemoryStore wrapped in an async context manager for compatibility.
+    """
+    store_manager = AsyncInMemoryStore()
+    yield await store_manager.__aenter__()

src/run_agent.py

@@ -0,0 +1,40 @@
+import asyncio
+from typing import cast
+from langsmith import uuid7
+
+from dotenv import load_dotenv
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.graph import MessagesState
+from langgraph.graph.state import CompiledStateGraph
+
+load_dotenv()
+
+from agents import DEFAULT_AGENT, get_agent  # noqa: E402
+
+# The default agent uses StateGraph.compile() which returns CompiledStateGraph
+agent = cast(CompiledStateGraph, get_agent(DEFAULT_AGENT))
+
+
+async def main() -> None:
+    inputs: MessagesState = {
+        "messages": [HumanMessage("Find me a recipe for chocolate chip cookies")]
+    }
+    result = await agent.ainvoke(
+        input=inputs,
+        config=RunnableConfig(configurable={"thread_id": uuid7()}),
+    )
+    result["messages"][-1].pretty_print()
+
+    # Draw the agent graph as png
+    # requires:
+    # brew install graphviz
+    # export CFLAGS="-I $(brew --prefix graphviz)/include"
+    # export LDFLAGS="-L $(brew --prefix graphviz)/lib"
+    # pip install pygraphviz
+    #
+    # agent.get_graph().draw_png("agent_diagram.png")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/run_service.py

@@ -0,0 +1,37 @@
+import asyncio
+import logging
+import sys
+
+import uvicorn
+from dotenv import load_dotenv
+
+from core import settings
+
+load_dotenv()
+
+if __name__ == "__main__":
+    root_logger = logging.getLogger()
+    if root_logger.handlers:
+        print(
+            f"Warning: Root logger already has {len(root_logger.handlers)} handler(s) configured. "
+            f"basicConfig() will be ignored. Current level: {logging.getLevelName(root_logger.level)}"
+        )
+
+    logging.basicConfig(level=settings.LOG_LEVEL.to_logging_level())
+    # Set Compatible event loop policy on Windows Systems.
+    # On Windows systems, the default ProactorEventLoop can cause issues with
+    # certain async database drivers like psycopg (PostgreSQL driver).
+    # The WindowsSelectorEventLoopPolicy provides better compatibility and prevents
+    # "RuntimeError: Event loop is closed" errors when working with database connections.
+    # This needs to be set before running the application server.
+    # Refer to the documentation for more information.
+    # https://www.psycopg.org/psycopg3/docs/advanced/async.html#asynchronous-operations
+    if sys.platform == "win32":
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+    uvicorn.run(
+        "service:app",
+        host=settings.HOST,
+        port=settings.PORT,
+        reload=settings.is_dev(),
+        timeout_graceful_shutdown=settings.GRACEFUL_SHUTDOWN_TIMEOUT,
+    )

src/schema/__init__.py

@@ -0,0 +1,25 @@
+from schema.models import AllModelEnum
+from schema.schema import (
+    AgentInfo,
+    ChatHistory,
+    ChatHistoryInput,
+    ChatMessage,
+    Feedback,
+    FeedbackResponse,
+    ServiceMetadata,
+    StreamInput,
+    UserInput,
+)
+
+__all__ = [
+    "AgentInfo",
+    "AllModelEnum",
+    "UserInput",
+    "ChatMessage",
+    "ServiceMetadata",
+    "StreamInput",
+    "Feedback",
+    "FeedbackResponse",
+    "ChatHistoryInput",
+    "ChatHistory",
+]

src/schema/models.py

@@ -0,0 +1,165 @@
+from enum import StrEnum, auto
+from typing import TypeAlias
+
+
+class Provider(StrEnum):
+    OPENAI = auto()
+    OPENAI_COMPATIBLE = auto()
+    AZURE_OPENAI = auto()
+    DEEPSEEK = auto()
+    ANTHROPIC = auto()
+    GOOGLE = auto()
+    VERTEXAI = auto()
+    GROQ = auto()
+    AWS = auto()
+    OLLAMA = auto()
+    OPENROUTER = auto()
+    FAKE = auto()
+
+
+class OpenAIModelName(StrEnum):
+    """https://platform.openai.com/docs/models/gpt-4o"""
+
+    GPT_5_NANO = "gpt-5-nano"
+    GPT_5_MINI = "gpt-5-mini"
+    GPT_5_1 = "gpt-5.1"
+
+
+class AzureOpenAIModelName(StrEnum):
+    """Azure OpenAI model names"""
+
+    AZURE_GPT_4O = "azure-gpt-4o"
+    AZURE_GPT_4O_MINI = "azure-gpt-4o-mini"
+
+
+class OpenAIEmbeddingModelName(StrEnum):
+    """https://platform.openai.com/docs/guides/embeddings"""
+
+    TEXT_EMBEDDING_3_SMALL = "text-embedding-3-small"
+    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
+    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"
+
+
+class DeepseekModelName(StrEnum):
+    """https://api-docs.deepseek.com/quick_start/pricing"""
+
+    DEEPSEEK_CHAT = "deepseek-chat"
+
+
+class AnthropicModelName(StrEnum):
+    """https://docs.anthropic.com/en/docs/about-claude/models#model-names"""
+
+    HAIKU_45 = "claude-haiku-4-5"
+    SONNET_45 = "claude-sonnet-4-5"
+
+
+class GoogleModelName(StrEnum):
+    """https://ai.google.dev/gemini-api/docs/models/gemini"""
+
+    GEMINI_15_PRO = "gemini-1.5-pro"
+    GEMINI_20_FLASH = "gemini-2.0-flash"
+    GEMINI_20_FLASH_LITE = "gemini-2.0-flash-lite"
+    GEMINI_25_FLASH = "gemini-2.5-flash"
+    GEMINI_25_PRO = "gemini-2.5-pro"
+    GEMINI_30_PRO = "gemini-3-pro-preview"
+
+
+class GoogleEmbeddingModelName(StrEnum):
+    """https://ai.google.dev/gemini-api/docs/models/gemini#text-embedding"""
+
+    TEXT_EMBEDDING_004 = "text-embedding-004"
+
+
+class VertexAIModelName(StrEnum):
+    """https://cloud.google.com/vertex-ai/generative-ai/docs/models"""
+
+    GEMINI_15_PRO = "gemini-1.5-pro"
+    GEMINI_20_FLASH = "gemini-2.0-flash"
+    GEMINI_20_FLASH_LITE = "models/gemini-2.0-flash-lite"
+    GEMINI_25_FLASH = "models/gemini-2.5-flash"
+    GEMINI_25_PRO = "gemini-2.5-pro"
+    GEMINI_30_PRO = "gemini-3-pro-preview"
+
+
+class GroqModelName(StrEnum):
+    """https://console.groq.com/docs/models"""
+
+    LLAMA_GUARD_4_12B = "meta-llama/llama-guard-4-12b"
+    LLAMA_31_8B_INSTANT = "llama-3.1-8b-instant"
+    LLAMA_33_70B_VERSATILE = "llama-3.3-70b-versatile"
+    LLAMA_4_MAVERICK_17B_128E = "meta-llama/llama-4-maverick-17b-128e-instruct"
+    LLAMA_4_SCOUT_17B_16E = "meta-llama/llama-4-scout-17b-16e-instruct"
+    LLAMA_PROMPT_GUARD_2_22M = "meta-llama/llama-prompt-guard-2-22m"
+    LLAMA_PROMPT_GUARD_2_86M = "meta-llama/llama-prompt-guard-2-86m"
+    OPENAI_GPT_OSS_120B = "openai/gpt-oss-120b"
+    OPENAI_GPT_OSS_20B = "openai/gpt-oss-20b"
+    OPENAI_GPT_OSS_SAFEGUARD_20B = "openai/gpt-oss-safeguard-20b"
+    GROQ_COMPOUND = "groq/compound"
+    GROQ_COMPOUND_MINI = "groq/compound-mini"
+    QWEN_3_32B = "qwen/qwen3-32b"
+    KIMI_K2_INSTRUCT = "moonshotai/kimi-k2-instruct"
+    KIMI_K2_INSTRUCT_0905 = "moonshotai/kimi-k2-instruct-0905"
+    ORPHEUS_ARABIC_SAUDI = "canopylabs/orpheus-arabic-saudi"
+    ORPHEUS_V1_ENGLISH = "canopylabs/orpheus-v1-english"
+    WHISPER_LARGE_V3 = "whisper-large-v3"
+    WHISPER_LARGE_V3_TURBO = "whisper-large-v3-turbo"
+    ALLAM_2_7B = "allam-2-7b"
+
+class AWSModelName(StrEnum):
+    """https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html"""
+
+    BEDROCK_HAIKU = "bedrock-3.5-haiku"
+    BEDROCK_SONNET = "bedrock-3.5-sonnet"
+
+
+class OllamaModelName(StrEnum):
+    """https://ollama.com/search"""
+
+    OLLAMA_GENERIC = "ollama"
+
+
+class OllamaEmbeddingModelName(StrEnum):
+    """Common Ollama embedding models"""
+
+    NOMIC_EMBED_TEXT = "nomic-embed-text"
+    ALL_MINILM = "all-minilm"
+
+
+class OpenRouterModelName(StrEnum):
+    """https://openrouter.ai/models"""
+
+    GEMINI_25_FLASH = "google/gemini-2.5-flash"
+
+
+class OpenAICompatibleName(StrEnum):
+    """https://platform.openai.com/docs/guides/text-generation"""
+
+    OPENAI_COMPATIBLE = "openai-compatible"
+
+
+class FakeModelName(StrEnum):
+    """Fake model for testing."""
+
+    FAKE = "fake"
+
+
+AllModelEnum: TypeAlias = (
+    OpenAIModelName
+    | OpenAICompatibleName
+    | AzureOpenAIModelName
+    | DeepseekModelName
+    | AnthropicModelName
+    | GoogleModelName
+    | VertexAIModelName
+    | GroqModelName
+    | AWSModelName
+    | OllamaModelName
+    | OpenRouterModelName
+    | FakeModelName
+)
+
+AllEmbeddingModelEnum: TypeAlias = (
+    OpenAIEmbeddingModelName
+    | GoogleEmbeddingModelName
+    | OllamaEmbeddingModelName
+)

src/schema/schema.py

@@ -0,0 +1,175 @@
+from typing import Any, Literal, NotRequired
+
+from pydantic import BaseModel, Field, SerializeAsAny
+from typing_extensions import TypedDict
+
+from schema.models import AllModelEnum, AnthropicModelName, OpenAIModelName
+
+
+class AgentInfo(BaseModel):
+    """Info about an available agent."""
+
+    key: str = Field(
+        description="Agent key.",
+        examples=["research-assistant"],
+    )
+    description: str = Field(
+        description="Description of the agent.",
+        examples=["A research assistant for generating research papers."],
+    )
+
+
+class ServiceMetadata(BaseModel):
+    """Metadata about the service including available agents and models."""
+
+    agents: list[AgentInfo] = Field(
+        description="List of available agents.",
+    )
+    models: list[AllModelEnum] = Field(
+        description="List of available LLMs.",
+    )
+    default_agent: str = Field(
+        description="Default agent used when none is specified.",
+        examples=["research-assistant"],
+    )
+    default_model: AllModelEnum = Field(
+        description="Default model used when none is specified.",
+    )
+
+
+class UserInput(BaseModel):
+    """Basic user input for the agent."""
+
+    message: str = Field(
+        description="User input to the agent.",
+        examples=["What is the weather in Tokyo?"],
+    )
+    model: SerializeAsAny[AllModelEnum] | None = Field(
+        title="Model",
+        description="LLM Model to use for the agent. Defaults to the default model set in the settings of the service.",
+        default=None,
+        examples=[OpenAIModelName.GPT_5_NANO, AnthropicModelName.HAIKU_45],
+    )
+    thread_id: str | None = Field(
+        description="Thread ID to persist and continue a multi-turn conversation.",
+        default=None,
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    user_id: str | None = Field(
+        description="User ID to persist and continue a conversation across multiple threads.",
+        default=None,
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    agent_config: dict[str, Any] = Field(
+        description="Additional configuration to pass through to the agent",
+        default={},
+        examples=[{"spicy_level": 0.8}],
+    )
+
+
+class StreamInput(UserInput):
+    """User input for streaming the agent's response."""
+
+    stream_tokens: bool = Field(
+        description="Whether to stream LLM tokens to the client.",
+        default=True,
+    )
+
+
+class ToolCall(TypedDict):
+    """Represents a request to call a tool."""
+
+    name: str
+    """The name of the tool to be called."""
+    args: dict[str, Any]
+    """The arguments to the tool call."""
+    id: str | None
+    """An identifier associated with the tool call."""
+    type: NotRequired[Literal["tool_call"]]
+
+
+class ChatMessage(BaseModel):
+    """Message in a chat."""
+
+    type: Literal["human", "ai", "tool", "custom"] = Field(
+        description="Role of the message.",
+        examples=["human", "ai", "tool", "custom"],
+    )
+    content: str = Field(
+        description="Content of the message.",
+        examples=["Hello, world!"],
+    )
+    tool_calls: list[ToolCall] = Field(
+        description="Tool calls in the message.",
+        default=[],
+    )
+    tool_call_id: str | None = Field(
+        description="Tool call that this message is responding to.",
+        default=None,
+        examples=["call_Jja7J89XsjrOLA5r!MEOW!SL"],
+    )
+    run_id: str | None = Field(
+        description="Run ID of the message.",
+        default=None,
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    response_metadata: dict[str, Any] = Field(
+        description="Response metadata. For example: response headers, logprobs, token counts.",
+        default={},
+    )
+    custom_data: dict[str, Any] = Field(
+        description="Custom message data.",
+        default={},
+    )
+
+    def pretty_repr(self) -> str:
+        """Get a pretty representation of the message."""
+        base_title = self.type.title() + " Message"
+        padded = " " + base_title + " "
+        sep_len = (80 - len(padded)) // 2
+        sep = "=" * sep_len
+        second_sep = sep + "=" if len(padded) % 2 else sep
+        title = f"{sep}{padded}{second_sep}"
+        return f"{title}\n\n{self.content}"
+
+    def pretty_print(self) -> None:
+        print(self.pretty_repr())  # noqa: T201
+
+
+class Feedback(BaseModel):  # type: ignore[no-redef]
+    """Feedback for a run, to record to LangSmith."""
+
+    run_id: str = Field(
+        description="Run ID to record feedback for.",
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    key: str = Field(
+        description="Feedback key.",
+        examples=["human-feedback-stars"],
+    )
+    score: float = Field(
+        description="Feedback score.",
+        examples=[0.8],
+    )
+    kwargs: dict[str, Any] = Field(
+        description="Additional feedback kwargs, passed to LangSmith.",
+        default={},
+        examples=[{"comment": "In-line human feedback"}],
+    )
+
+
+class FeedbackResponse(BaseModel):
+    status: Literal["success"] = "success"
+
+
+class ChatHistoryInput(BaseModel):
+    """Input for retrieving chat history."""
+
+    thread_id: str = Field(
+        description="Thread ID to persist and continue a multi-turn conversation.",
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+
+
+class ChatHistory(BaseModel):
+    messages: list[ChatMessage]

src/schema/task_data.py

@@ -0,0 +1,74 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class TaskData(BaseModel):
+    name: str | None = Field(
+        description="Name of the task.", default=None, examples=["Check input safety"]
+    )
+    run_id: str = Field(
+        description="ID of the task run to pair state updates to.",
+        default="",
+        examples=["847c6285-8fc9-4560-a83f-4e6285809254"],
+    )
+    state: Literal["new", "running", "complete"] | None = Field(
+        description="Current state of given task instance.",
+        default=None,
+        examples=["running"],
+    )
+    result: Literal["success", "error"] | None = Field(
+        description="Result of given task instance.",
+        default=None,
+        examples=["running"],
+    )
+    data: dict[str, Any] = Field(
+        description="Additional data generated by the task.",
+        default={},
+    )
+
+    def completed(self) -> bool:
+        return self.state == "complete"
+
+    def completed_with_error(self) -> bool:
+        return self.state == "complete" and self.result == "error"
+
+
+class TaskDataStatus:
+    def __init__(self) -> None:
+        import streamlit as st
+
+        self.status = st.status("")
+        self.current_task_data: dict[str, TaskData] = {}
+
+    def add_and_draw_task_data(self, task_data: TaskData) -> None:
+        status = self.status
+        status_str = f"Task **{task_data.name}** "
+        match task_data.state:
+            case "new":
+                status_str += "has :blue[started]. Input:"
+            case "running":
+                status_str += "wrote:"
+            case "complete":
+                if task_data.result == "success":
+                    status_str += ":green[completed successfully]. Output:"
+                else:
+                    status_str += ":red[ended with error]. Output:"
+        status.write(status_str)
+        status.write(task_data.data)
+        status.write("---")
+        if task_data.run_id not in self.current_task_data:
+            # Status label always shows the last newly started task
+            status.update(label=f"""Task: {task_data.name}""")
+        self.current_task_data[task_data.run_id] = task_data
+        if all(entry.completed() for entry in self.current_task_data.values()):
+            # Status is "error" if any task has errored
+            if any(entry.completed_with_error() for entry in self.current_task_data.values()):
+                state = "error"
+            # Status is "complete" if all tasks have completed successfully
+            else:
+                state = "complete"
+        # Status is "running" until all tasks have completed
+        else:
+            state = "running"
+        status.update(state=state)  # type: ignore[arg-type]

src/scripts/create_chroma_db.py

@@ -0,0 +1,83 @@
+import os
+import shutil
+
+from dotenv import load_dotenv
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader
+from langchain_openai import OpenAIEmbeddings
+
+# Load environment variables from the .env file
+load_dotenv()
+
+
+def create_chroma_db(
+    folder_path: str,
+    db_name: str = "./chroma_db",
+    delete_chroma_db: bool = True,
+    chunk_size: int = 2000,
+    overlap: int = 500,
+):
+    embeddings = OpenAIEmbeddings(api_key=os.environ["OPENAI_API_KEY"])
+
+    # Initialize Chroma vector store
+    if delete_chroma_db and os.path.exists(db_name):
+        shutil.rmtree(db_name)
+        print(f"Deleted existing database at {db_name}")
+
+    chroma = Chroma(
+        embedding_function=embeddings,
+        persist_directory=f"./{db_name}",
+    )
+
+    # Initialize text splitter
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=overlap)
+
+    # Iterate over files in the folder
+    for filename in os.listdir(folder_path):
+        file_path = os.path.join(folder_path, filename)
+
+        # Load document based on file extension
+        # Add more loaders if required, i.e. JSONLoader, TxtLoader, etc.
+        if filename.endswith(".pdf"):
+            loader = PyPDFLoader(file_path)
+        elif filename.endswith(".docx"):
+            loader = Docx2txtLoader(file_path)
+        else:
+            continue  # Skip unsupported file types
+
+        # Load and split document into chunks
+        document = loader.load()
+        chunks = text_splitter.split_documents(document)
+
+        # Add chunks to Chroma vector store
+        for chunk in chunks:
+            chunk_id = chroma.add_documents([chunk])
+            if chunk_id:
+                print(f"Chunk added with ID: {chunk_id}")
+            else:
+                print("Failed to add chunk")
+
+        print(f"Document {filename} added to database.")
+
+    print(f"Vector database created and saved in {db_name}.")
+    return chroma
+
+
+if __name__ == "__main__":
+    # Path to the folder containing the documents
+    folder_path = "./data"
+
+    # Create the Chroma database
+    chroma = create_chroma_db(folder_path=folder_path)
+
+    # Create retriever from the Chroma database
+    retriever = chroma.as_retriever(search_kwargs={"k": 3})
+
+    # Perform a similarity search
+    query = "What's my company's mission and values"
+    similar_docs = retriever.invoke(query)
+
+    # Display results
+    for i, doc in enumerate(similar_docs, start=1):
+        print(f"\n🔹 Result {i}:\n{doc.page_content}\nTags: {doc.metadata.get('source', [])}")

src/scripts/load_portfolio.py

@@ -0,0 +1,25 @@
+import argparse
+import logging
+from dotenv import load_dotenv
+
+from core.settings import settings
+from scripts.portfolio.portfolio_ingestion import PortfolioIngest
+
+logging.basicConfig(
+    level=logging.INFO,  # Use INFO level to see all sync progress logs
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+
+load_dotenv()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Synchronize portfolio data from Notion")
+    parser.add_argument(
+        "--since",
+        type=str,
+        help="ISO 8601 date to sync from (e.g. 2024-01-01T00:00:00.000Z). If not provided, uses last sync date."
+    )
+    args = parser.parse_args()
+    
+    orchestrator = PortfolioIngest()
+    orchestrator.sync(args.since)

src/scripts/portfolio/document.py

@@ -0,0 +1,129 @@
+import json
+import hashlib
+import time
+import re
+from typing import Optional, List, Tuple
+
+from langsmith import uuid7
+from langchain_core.documents import Document
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableConfig
+from langchain_core.documents import Document
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+from core.llm import get_model
+from core.settings import settings
+from scripts.portfolio.prompt import PORTFOLIO_INGESTION_SYSTEM_PROMPT
+
+class DocumentChunker:
+    """Service for splitting documents into chunks."""
+    
+    def __init__(self, chunk_size: int = 1500, chunk_overlap: int = 200):
+        self.text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            is_separator_regex=True,
+        )
+        print(f"DEBUG: Initialized DocumentChunker with chunk_size={chunk_size}, overlap={chunk_overlap}")
+    
+    def chunk_document(self, doc: Document, base_id: str, content_hash: str) -> List[Tuple[Document, str]]:
+        """
+        Splits a document into chunks and prepares them for storage.
+        
+        Args:
+            doc: The document to chunk
+            base_id: The base document ID
+            content_hash: The content hash for change detection
+            
+        Returns:
+            List of tuples (chunk_document, chunk_id)
+        """
+        chunks = self.text_splitter.split_documents([doc])
+        chunked_docs = []
+        
+        for idx, chunk in enumerate(chunks):
+            chunk_id = f"{base_id}_chunk_{idx}"
+            chunk.metadata["content_hash"] = content_hash
+            chunk.metadata["base_id"] = base_id
+            chunked_docs.append((chunk, chunk_id))
+        
+        print(f"DEBUG: Split document {base_id} into {len(chunked_docs)} chunks")
+        return chunked_docs
+
+
+class DocumentEnricher:
+    """Service for enriching documents using LLM with generalized retry logic."""
+    
+    def __init__(self):
+        self.llm = get_model(settings.DEFAULT_MODEL)
+        self.enrich_prompt = ChatPromptTemplate.from_messages([
+            ("system", PORTFOLIO_INGESTION_SYSTEM_PROMPT),
+            ("human", "Category: {category}\n\nMetadata:\n{metadata}\n\nContent:\n{content}")
+        ])
+        print(f"INFO: Initialized DocumentEnricher with {settings.DEFAULT_MODEL}")
+    
+    def enrich(self, doc: Document, category: str, max_retries: int = 5) -> Tuple[Optional[Document], str, str]:
+        pid = str(doc.metadata.get("id", uuid7()))
+        title = doc.metadata.get("Title", "Untitled")
+        
+        for attempt in range(max_retries):
+            try:
+                if attempt > 0:
+                    wait_time = min(2 ** attempt, 60)
+                    print(f"INFO: Retrying {title} (attempt {attempt + 1}/{max_retries}) in {wait_time}s...")
+                    time.sleep(wait_time)
+                else:
+                    print(f"INFO: Enriching document: {title} (PID: {pid})")
+                
+                res = self.llm.invoke(
+                    self.enrich_prompt.format_messages(
+                        category=category,
+                        metadata=json.dumps(doc.metadata, default=str),
+                        content=doc.page_content or "No content provided."
+                    ),
+                    config=RunnableConfig(run_id=uuid7())
+                )
+                
+                enriched_content = res.content.strip()
+                content_hash = hashlib.sha256(enriched_content.encode('utf-8')).hexdigest()
+                
+                enriched_doc = Document(
+                    page_content=enriched_content,
+                    metadata={
+                        **doc.metadata,
+                        "category": category,
+                        "content_hash": content_hash,
+                        "base_id": pid
+                    }
+                )
+                return enriched_doc, pid, content_hash
+                
+            except Exception as e:
+                error_msg = str(e).lower()
+                error_type = type(e).__name__.lower()
+                
+                # --- Rate Limit Detection ---
+                is_rate_limit = any(keyword in error_msg or keyword in error_type 
+                                   for keyword in ["429", "rate_limit", "rate limit", "too many requests", "throttled"])
+                
+                # --- Overloaded/Server Error Detection ---
+                is_server_error = any(keyword in error_msg 
+                                     for keyword in ["500", "502", "503", "overloaded", "unavailable", "deadline_exceeded"])
+
+                if is_rate_limit or is_server_error:
+                    wait_time = 5 # Default
+                    match = re.search(r'(?:try again in|retry after|wait)\s*([\d.]+)\s*s', error_msg)
+                    if match:
+                        wait_time = float(match.group(1)) + 1
+                    
+                    if attempt < max_retries - 1:
+                        print(f"WARN: API issue (Rate Limit/Overload) for {title}. Waiting {wait_time}s...")
+                        time.sleep(wait_time)
+                        continue
+                
+                # Non-retriable or final attempt failure
+                print(f"ERROR: Enrichment failed for {title}: {e}")
+                if attempt >= 1:
+                    return None, pid, ""
+
+        return None, pid, ""

src/scripts/portfolio/notion_loader.py

@@ -0,0 +1,68 @@
+import os
+import traceback
+from typing import List, Optional
+
+from langchain_community.document_loaders import NotionDBLoader
+from langchain_core.documents import Document
+
+NOTION_DB_MAP = {
+    "education": "NOTION_EDUCATION_ID",
+    "experience": "NOTION_EXPERIENCE_ID",
+    "projects": "NOTION_PROJECT_ID",
+    "testimonials": "NOTION_TESTIMONIAL_ID",
+    "blog": "NOTION_BLOG_ID",
+}
+class NotionLoader:
+    """Service for loading documents from Notion databases."""
+    
+    def __init__(self):
+        self.token = os.getenv("NOTION_TOKEN")
+        if not self.token:
+            print("WARNING: NOTION_TOKEN not found in environment")
+    
+    def load_category(self, category: str, since_date: Optional[str] = None) -> List[Document]:
+        """
+        Loads documents from a Notion database for a given category.
+        
+        Args:
+            category: The category name (must be in NOTION_DB_MAP)
+            since_date: Optional ISO 8601 date to filter documents updated after this date
+            
+        Returns:
+            List of documents that were updated after since_date
+        """
+        if category not in NOTION_DB_MAP:
+            print(f"WARNING: Unknown category: {category}")
+            return []
+        
+        env_key = NOTION_DB_MAP[category]
+        db_id = os.getenv(env_key)
+        
+        if not db_id:
+            print(f"WARNING: Notion database ID not found for category: {category}")
+            return []
+        
+        if not self.token:
+            print("ERROR: NOTION_TOKEN not available")
+            return []
+        
+        try:
+            print(f"INFO: Loading {category} documents from Notion...")
+            loader = NotionDBLoader(self.token, db_id)
+            all_docs = loader.load()
+            
+            if since_date:
+                valid_docs = [
+                    d for d in all_docs
+                    if d.metadata.get("updated", "") > since_date
+                ]
+                print(f"INFO: Found {len(valid_docs)} documents updated after {since_date} out of {len(all_docs)} total")
+            else:
+                valid_docs = all_docs
+                print(f"INFO: Loaded {len(valid_docs)} documents (no date filter)")
+            
+            return valid_docs
+        except Exception as e:
+            print(f"ERROR: Failed to load {category} from Notion: {e}")
+            traceback.print_exc()
+            return []

src/scripts/portfolio/portfolio_ingestion.py

@@ -0,0 +1,150 @@
+import traceback
+from typing import Optional, Tuple
+
+from langchain_community.vectorstores.utils import filter_complex_metadata
+
+from core.settings import settings
+from memory.postgres import load_pgvector_store
+from scripts.portfolio.document import DocumentChunker, DocumentEnricher
+from scripts.portfolio.notion_loader import NotionLoader
+from scripts.portfolio.vector_repository import VectorRepository
+NOTION_DB_MAP = {
+    "education": "NOTION_EDUCATION_ID",
+    "experience": "NOTION_EXPERIENCE_ID",
+    "projects": "NOTION_PROJECT_ID",
+    "testimonials": "NOTION_TESTIMONIAL_ID",
+    "blog": "NOTION_BLOG_ID",
+}
+
+class PortfolioIngest:
+    """Orchestrates the portfolio synchronization process."""
+    
+    def __init__(self):
+        self.enricher = DocumentEnricher()
+        self.chunker = DocumentChunker()
+        self.loader = NotionLoader()
+        self.repository = VectorRepository()
+        self.store = load_pgvector_store()
+        print(f"INFO: Initialized PortfolioIngest with collection: {settings.VECTOR_STORE_COLLECTION_NAME}")
+    
+    def ingest_category(self, category: str, since_date: str) -> Tuple[int, int, int, int]:
+        """
+        Synchronizes a single category of documents.
+        
+        Args:
+            category: The category to sync
+            since_date: ISO 8601 date to filter documents
+            
+        Returns:
+            Tuple of (chunks_deleted, chunks_updated, chunks_skipped, total_synced)
+        """
+        # Load documents from Notion
+        valid_docs = self.loader.load_category(category, since_date)
+        if not valid_docs:
+            print(f"INFO: No new updates for {category}")
+            return 0, 0, 0, 0
+        
+        # Enrich documents with LLM sequentially to avoid rate limits
+        print(f"INFO: Enriching {len(valid_docs)} documents with LLM...")
+        
+        enriched_docs = []
+        base_ids = []
+        content_hashes = []
+        
+        for doc in valid_docs:
+            enriched_doc, base_id, content_hash = self.enricher.enrich(doc, category)
+            if enriched_doc is not None:
+                enriched_docs.append(enriched_doc)
+                base_ids.append(base_id)
+                content_hashes.append(content_hash)
+        
+        if not enriched_docs:
+            print(f"INFO: No enriched documents for {category}")
+            return 0, 0, 0, 0
+        
+        # Batch fetch existing content hashes
+        print(f"INFO: Checking content hashes for {len(enriched_docs)} documents...")
+        existing_hashes = self.repository.batch_get_existing_content_hashes(base_ids)
+        print(f"DEBUG: Found {len([h for h in existing_hashes.values() if h is not None])} existing hashes")
+        
+        # Process documents and prepare for upsert
+        docs_to_upsert = []
+        ids_to_upsert = []
+        chunks_deleted = 0
+        chunks_updated = 0
+        chunks_skipped = 0
+        
+        for enriched_doc, base_id, new_hash in zip(enriched_docs, base_ids, content_hashes):
+            existing_hash = existing_hashes.get(base_id)
+            
+            if existing_hash is None:
+                # New document
+                print(f"INFO: New document {base_id}, adding all chunks...")
+                chunked = self.chunker.chunk_document(enriched_doc, base_id, new_hash)
+                for chunk, chunk_id in chunked:
+                    docs_to_upsert.append(chunk)
+                    ids_to_upsert.append(chunk_id)
+                chunks_updated += len(chunked)
+            elif existing_hash != new_hash:
+                # Content changed
+                print(f"INFO: Content changed for {base_id}, replacing chunks...")
+                old_chunk_ids = self.repository.get_existing_chunks(base_id)
+                if old_chunk_ids:
+                    self.store.delete(old_chunk_ids)
+                    chunks_deleted += len(old_chunk_ids)
+                    print(f"DEBUG: Deleted {len(old_chunk_ids)} old chunks for {base_id}")
+                
+                chunked = self.chunker.chunk_document(enriched_doc, base_id, new_hash)
+                for chunk, chunk_id in chunked:
+                    docs_to_upsert.append(chunk)
+                    ids_to_upsert.append(chunk_id)
+                chunks_updated += len(chunked)
+            else:
+                # Content unchanged
+                print(f"DEBUG: Content unchanged for {base_id}, skipping...")
+                chunks_skipped += 1
+        
+        # Upsert documents
+        total_synced = 0
+        if docs_to_upsert:
+            print(f"INFO: Upserting {len(docs_to_upsert)} chunks...")
+            self.store.add_documents(
+                filter_complex_metadata(docs_to_upsert),
+                ids=ids_to_upsert
+            )
+            total_synced = len(docs_to_upsert)
+            print(
+                f"INFO: Category {category}: Deleted {chunks_deleted} chunks, "
+                f"updated {chunks_updated} chunks, skipped {chunks_skipped} unchanged documents, "
+                f"total synced: {total_synced}"
+            )
+        
+        return chunks_deleted, chunks_updated, chunks_skipped, total_synced
+    
+    def sync(self, manual_date: Optional[str] = None) -> int:
+        """
+        Synchronizes all portfolio categories.
+        
+        Args:
+            manual_date: Optional ISO 8601 date to override last sync date
+            
+        Returns:
+            Total number of chunks synced
+        """
+        since_date = manual_date or self.repository.get_last_sync_date()
+        print(f"INFO: Starting sync for items modified after: {since_date}")
+        
+        total_synced = 0
+        
+        for category in NOTION_DB_MAP.keys():
+            try:
+                chunks_deleted, chunks_updated, chunks_skipped, synced = self.ingest_category(
+                    category, since_date
+                )
+                total_synced += synced
+            except Exception as e:
+                print(f"ERROR: Failed to sync category {category}: {e}")
+                traceback.print_exc()
+        
+        print(f"INFO: Sync complete. Total chunks synced: {total_synced}")
+        return total_synced

src/scripts/portfolio/prompt.py

@@ -0,0 +1,29 @@
+import os
+from dotenv import load_dotenv
+from datetime import datetime
+
+load_dotenv()
+
+OWNER = os.getenv("OWNER", "Anuj Joshi")
+CURRENT_DATE = datetime.now().strftime("%Y-%m-%d")
+
+PORTFOLIO_INGESTION_SYSTEM_PROMPT = f"""
+SYSTEM ROLE: You are the Principal Technical Documentation Engineer for {OWNER}.
+
+Your goal is to transform raw project data into high-density retrieval-optimized technical documentation for {OWNER}.
+This documentation is the primary source for a RAG (Retrieval-Augmented Generation) system used by Staff Engineers and Technical Recruiters.
+
+General Behavior
+- Do include a header block that summarizes the content of the document in a concise and catchy manner.
+- Do include thumbnail/links if provided in the header block if available.
+- Do include other links in other sections in proper markdown format if available and relevant.
+- Use "Information-Dense Sentences." Every sentence must provide a new fact.
+- Avoid pronouns, repeat the subject or {OWNER}'s name to ensure chunks remain context-aware when retrieved in isolation.
+- If data is missing, do not hallucinate and omit the specific metric.
+- Do NOT use vague impact words.
+- Do NOT restate the same idea in different words.
+- Do NOT invent metrics, scale, or outcomes.
+- Do NOT include emojis or stylistic symbols.
+- No repetitive framing, long paragraphs, or storytelling.
+- Optimize for semantic search and chunk retrieval.
+"""