RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

v1.0.0 Release Candidate (#95) 

* chore: Update versioning for release

* fix: remove db creation at start and sync migrations and models

* fix: Checkpoint changing metamessage schema

* chore: linter fixes

* fix: session cloning working

* Hybrid long-term memory (#92)

* Add TOM method switching

* Add system prompt and note on format

* Add persistence tweaks

* Specify format for each section of user representation

* Parse XML tags before saving representation metamessage

* Clean up

* Use Claude 3.5 Haiku and refine prompt

* Simplify message processing

* chore: update token limit on dialectic and model for deriver

* Add embedding-based long-term fact retrieval

* Fix bug preventing new documents from being created

* Use multiple queries + tweak prompt

* Fix collection name bug + add duplicate removal

* First implementation of on-demand user rep generation

* WIP debug on-demand user rep changes

* Fixed representations not being stored & deriver issue

* Some speed improvements

* Play with number of facts / queries

* WIP prompt caching for Claude

* WIP fix anthropic caching

* Anthropic prompt caching working but messages too short

* Use Cerebras for small inferences

* Make dialectic responses 1000 tokens max

* Make user representation generation model a constant

* Use llama 3.1 8b for query generation

* Update env template

* Add crud.get_or_create_protected_collection

* rabbit comments

* Fix linter issues

* Add Cerebras to stream router method

* Better handling of default-empty string args

* Change prints to debug logs

* Add error handling to TOM inference

* Handle missing/empty client in model responses

* Handle no messages case in get_chat_history

* Fix indent

* Add error handling to single_prompt methods

* Fix get_or_create_user_protected_collection

* Simplify openAI-compatible model client instantiation

* Remove health endpoint

* Remove LocalEmbeddingStore

* Change prints to debug logs

* Change sentry track

* Code review changes

* Add README to ToM module

* Switch to Groq

* Fix inconsistent openai compatible provider list in stream()

* Update env template to include Groq variables

* Add model_client tests

* fix: Fix unit tests

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>

* add scoped API keys (#91)

* add AUTH_JWT_SECRET and ADMIN_KEY, use in security middleware (TODO granular keys)

* WIP: convert all API paths to use scoped keys

* add basic unit tests for API keys, ruff formatting

* MVP of route using JWT for payload

* add get_user_from_token

* add key table to postgres, use it to enable key revocation

* add key revocation pt 2 -- fix order of param checks

* finish convenience routes that assume params from JWT

* add tests for key API

* get_keys

* add secrets utility script, add key rotation, fill out tests

* add tiny cache as PoC

* nits, validations, etc

* only create keys table migration if necessary

* fix keys tests to always use auth

* tiny fix to make custom DATABASE_SCHEMA work

* review: add better docs, fix security issue with cache, clear db on rotation, and more

* remove rotation

* remove key database entirely

* Add `/all` path to get all apps (#94)

* add `/all` path for apps

* assert vector extension installed (need this for groudon)

* review

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>

* add scoped API keys (#91)

* add AUTH_JWT_SECRET and ADMIN_KEY, use in security middleware (TODO granular keys)

* WIP: convert all API paths to use scoped keys

* add basic unit tests for API keys, ruff formatting

* MVP of route using JWT for payload

* add get_user_from_token

* add key table to postgres, use it to enable key revocation

* add key revocation pt 2 -- fix order of param checks

* finish convenience routes that assume params from JWT

* add tests for key API

* get_keys

* add secrets utility script, add key rotation, fill out tests

* add tiny cache as PoC

* nits, validations, etc

* only create keys table migration if necessary

* fix keys tests to always use auth

* tiny fix to make custom DATABASE_SCHEMA work

* review: add better docs, fix security issue with cache, clear db on rotation, and more

* remove rotation

* remove key database entirely

* Add `/all` path to get all apps (#94)

* add `/all` path for apps

* assert vector extension installed (need this for groudon)

* review

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>

* chore: README and CHANGELOG updates

* add JWT expiry

* fix: Consolidate get methods with JWT token resolution

* chore: Add Annotation to Path, Query, and Body params

* chore: run ruff formatter

* chore: nits & add one exhaustive test of a query route

* fix: undo change to fly.toml

* fix: Langfuse tracing

* Consolidate Get Methods (#96)

* fix: Consolidate get methods with JWT token resolution

* chore: Add Annotation to Path, Query, and Body params

* chore: run ruff formatter

* chore: nits & add one exhaustive test of a query route

* fix: undo change to fly.toml

---------

Co-authored-by: dr-frmr <docterformer@protonmail.com>

* fix: dev-667 fix streaming endpoint

* fix: Anthropic Langfuse Tracing

* fix: add scripts folder to dockerfile

* fix: Remove redundant fields from pydantic schemas

* fix: Add deeper protection on reserved collection

* fix: Consolidate chat and stream methods

* docs: Update Mintlify API Reference and Changelog

* remove langchain guide, update architecture diagram

* honcho mcp server

* chore: Update .env template

* update discord, temporarily remove other guides

* Limit dialectic & deriver context usage with two-scale progressive summarization (#97)

* WIP two tiered summaries

* Move to process_item

* Save user rep metamessage even if no message_id

* Change number of messages per short summary

* Fix broken mock

* Remove prints

* chore: fix test

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>

* feat: Add Gemini Support, link facts to message, use 8b for dialectic fact queries

* chore: Styling

* chore: coderabbit nitpicks

* keep dialectic guide

* Add streaming guide

* Remove TODO from dialectic guide

* Fix JS snippets that referred to honcho singleton as client

* Add App explanation to architecture page

---------

Co-authored-by: Dani Balcells <18307962+danibalcells@users.noreply.github.com>
Co-authored-by: doria <93405247+dr-frmr@users.noreply.github.com>
Co-authored-by: dr-frmr <docterformer@protonmail.com>
Co-authored-by: vintro <vince@plasticlabs.ai>
Co-authored-by: Daniel Balcells <dbalcells@gmail.com>
This commit is contained in:
Vineeth Voruganti 2025-04-10 13:59:40 -04:00 committed by GitHub
parent c2a50241cc
commit 8588d36eb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
79 changed files with 26274 additions and 4930 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

43
.env.template
View File

@ -1,23 +1,44 @@
CONNECTION_URI=postgresql+psycopg://testuser:testpwd@localhost:5432/honcho # sample for local database
# CONNECTION_URI=postgresql+psycopg://testuser:testpwd@database:5432/honcho # sample for docker-compose database
OPENAI_API_KEY= # Used for vector embeddings
ANTHROPIC_API_KEY= # Used for the deriver and dialectic API
# Logging
SENTRY_ENABLED=false # Set to true to enable Sentry logging and tracing
# Use something unique here if you want to share a database with other projects.
# Leave blank for public. Make sure to avoid `-` in name.
DATABASE_SCHEMA=
# Auth
USE_AUTH_SERVICE=false
SECRET_KEY=
# Set to true to enable API authorization. Blank is equivalent to false.
USE_AUTH=false
# Required if USE_AUTH is true. Generate with scripts/generate_jwt_secret.py
AUTH_JWT_SECRET=
# These are included for convenience in local testing if you want to quickly swap out providers
# but are not actually used by Honcho
OPENROUTER_API_KEY=
OPENROUTER_BASE_URL=https://openrouter.ai/api/v1
CEREBRAS_API_KEY=
CEREBRAS_BASE_URL=https://api.cerebras.ai/v1
GROQ_API_KEY=
GROQ_BASE_URL=https://api.groq.com/openai/v1
# These are the ones that are actually used by the model client
OPENAI_COMPATIBLE_BASE_URL=
OPENAI_COMPATIBLE_API_KEY=
# Sentry
SENTRY_ENABLED=false
SENTRY_DSN=
# Deriver
DERIVER_WORKERS=1
TOM_METHOD="single_prompt"
USER_REPRESENTATION_METHOD="single_prompt"
# Anthropic
ANTHROPIC_API_KEY=
OPENAI_API_KEY=
# Langfuse
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
LANGFUSE_HOST=https://us.cloud.langfuse.com
# set logger level
LOG_LEVEL=INFO

18
.github/workflows/fly-deploy.yml vendored Normal file
View File

@ -0,0 +1,18 @@
# See https://fly.io/docs/app-guides/continuous-deployment-with-github-actions/
name: Fly Deploy
on:
push:
branches:
- main
jobs:
deploy:
name: Deploy app
runs-on: ubuntu-latest
concurrency: deploy-group # optional: ensure only one action runs at a time
steps:
- uses: actions/checkout@v4
- uses: superfly/flyctl-actions/setup-flyctl@master
- run: flyctl deploy --remote-only
env:
FLY_API_TOKEN: ${{ secrets.FLY_API_TOKEN }}

2
.github/workflows/unittest.yml vendored
View File

@ -47,7 +47,7 @@ jobs:
run: uv run pytest -x
env:
CONNECTION_URI: postgresql+psycopg://postgres:postgres@localhost:5432/test_db
USE_AUTH_SERVICE: false
USE_AUTH: false
SENTRY_ENABLED: false
OPENTELEMETRY_ENABLED: false
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}

22
CHANGELOG.md
View File

@ -5,14 +5,30 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](http://keepachangelog.com/)
and this project adheres to [Semantic Versioning](http://semver.org/).
## [1.0.0]
### Added
- JWT based API authentication
- Configurable logging
- Consolidated LLM Inference via `ModelClient` class
- Dynamic logging configurable via environment variables
### Changed
- Deriver & Dialectic API to use Hybrid Memory Architecture
- Metamessages are not strictly tied to a message
- Database provisioning is a separate script instead of happening on startup
- Consolidated `session/chat` and `session/chat/stream` endpoints
## [0.0.16]
## Added
### Added
- Detailed custom exceptions for better error handling
- CLAUDE.md for claude code
## Changed
### Changed
- Deriver to use a new cognitive architecture that only updates on user messages
and updates user representation to apply more confidence scores to its known
@ -22,7 +38,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
- SQLAlchemy echo changed to false by default, can be enabled with SQL_DEBUG
environment flag
## Fixed
### Fixed
- Self-hosting documentation and README to mention `uv` instead of `poetry`

16
CONTRIBUTING.md
View File

@ -91,10 +91,24 @@ that they are disabled you can verify the following environment variables are
set to false.
```env
USE_AUTH_SERVICE=false
USE_AUTH=false
SENTRY_ENABLED=false
```
If you set `USE_AUTH` to true you will need to generate a JWT secret. You can
do this with the following command:
```bash
python scripts/generate_jwt_secret.py
```
This will generate a JWT secret and print it to the console. You can then set
the `AUTH_JWT_SECRET` environment variable. This is required for `USE_AUTH`.
```env
AUTH_JWT_SECRET=<generated_secret>
```
5. Launch the API
With the dependencies installed, a database setup and enabled with `pgvector`,

1
Dockerfile
View File

@ -35,6 +35,7 @@ ENV PATH="/app/.venv/bin:$PATH"
COPY --chown=app:app src/ /app/src/
COPY --chown=app:app migrations/ /app/migrations/
COPY --chown=app:app scripts/ /app/scripts/
COPY --chown=app:app alembic.ini /app/alembic.ini
EXPOSE 8000

34
README.md
View File

@ -1,6 +1,6 @@
# 🫡 Honcho
![Static Badge](https://img.shields.io/badge/Version-0.0.16-blue)
![Static Badge](https://img.shields.io/badge/Version-1.0.0-blue)
[![Discord](https://img.shields.io/discord/1016845111637839922?style=flat&logo=discord&logoColor=23ffffff&label=Plastic%20Labs&labelColor=235865F2)](https://discord.gg/plasticlabs)
[![arXiv](https://img.shields.io/badge/arXiv-2310.06983-b31b1b.svg)](https://arxiv.org/abs/2310.06983)
![GitHub License](https://img.shields.io/github/license/plastic-labs/honcho)
@ -75,10 +75,10 @@ Below is a mapping of the different primitives.
Apps
└── Users
├── Sessions
│ ├── Messages
│ └── Metamessages
└── Collections
└── Documents
│ └── Messages
├── Collections
│ └── Documents
└── Metamessages
```
Users familiar with APIs such as the OpenAI Assistants API will be familiar with
@ -105,18 +105,6 @@ The `Session` object represents a set of interactions a `User` has with an
The `Message` represents an atomic interaction of a `User` in a `Session`.
`Message`s are labed as either a `User` or AI message.
#### Metamessages
A `Metamessage` is similar to a `Message` with different use case. They are
meant to be used to store intermediate inference from AI assistants or other
derived information that is separate from the main `User` `App` interaction
loop. For complicated prompting architectures like [metacognitive prompting](https://arxiv.org/abs/2310.06983)
metamessages can store thought and reflection steps along with having developer
information such as logs.
Each `Metamessage` is associated with a `Message`. The convention we recommend
is to attach a `Metamessage` to the `Message` it was derived from or based on.
#### Collections
At a high level a `Collection` is a named group of `Documents`. Developers
@ -132,6 +120,18 @@ PDF files, and more.
As stated before a `Document` is vector embedded data stored in a `Collection`.
#### Metamessages
A `Metamessage` is similar to a `Message` with different use case. They are
meant to be used to store intermediate inference from AI assistants or other
derived information that is separate from the main `User` `App` interaction
loop. For complicated prompting architectures like [metacognitive prompting](https://arxiv.org/abs/2310.06983)
metamessages can store thought and reflection steps along with having developer
information such as logs.
Each `Metamessage` is associated with a `User` with the ability to optionally
tie to a `Session` and a `Message`.
### Insights
The Insight functionality of Honcho is built on top of the Storage service. As

3
docs/api-reference/endpoint/apps/get-all-apps.mdx Normal file
View File

@ -0,0 +1,3 @@
---
openapi: post /v1/apps/list
---

4
docs/api-reference/endpoint/apps/get-app.mdx
View File

@ -1,3 +1,3 @@
---
openapi: get /v1/apps/{app_id}
---
openapi: get /v1/apps
---

3
docs/api-reference/endpoint/collections/get-collection-by-id.mdx
View File

@ -1,3 +0,0 @@
---
openapi: get /v1/apps/{app_id}/users/{user_id}/collections/{collection_id}
---

3
docs/api-reference/endpoint/collections/get-collection.mdx Normal file
View File

@ -0,0 +1,3 @@
---
openapi: get /v1/apps/{app_id}/users/{user_id}/collections
---

3
docs/api-reference/endpoint/keys/create-key.mdx Normal file
View File

@ -0,0 +1,3 @@
---
openapi: post /v1/keys
---

4
docs/api-reference/endpoint/metamessages/create-metamessage.mdx
View File

@ -1,3 +1,3 @@
---
openapi: post /v1/apps/{app_id}/users/{user_id}/sessions/{session_id}/metamessages
---
openapi: post /v1/apps/{app_id}/users/{user_id}/metamessages
---

4
docs/api-reference/endpoint/metamessages/get-metamessage.mdx
View File

@ -1,3 +1,3 @@
---
openapi: get /v1/apps/{app_id}/users/{user_id}/sessions/{session_id}/metamessages/{metamessage_id}
---
openapi: get /v1/apps/{app_id}/users/{user_id}/metamessages/{metamessage_id}
---

3
docs/api-reference/endpoint/metamessages/get-metamessages-by-user.mdx
View File

@ -1,3 +0,0 @@
---
openapi: post /v1/apps/{app_id}/users/{user_id}/metamessages/list
---

4
docs/api-reference/endpoint/metamessages/get-metamessages.mdx
View File

@ -1,3 +1,3 @@
---
openapi: post /v1/apps/{app_id}/users/{user_id}/sessions/{session_id}/metamessages/list
---
openapi: post /v1/apps/{app_id}/users/{user_id}/metamessages/list
---

4
docs/api-reference/endpoint/metamessages/update-metamessage.mdx
View File

@ -1,3 +1,3 @@
---
openapi: put /v1/apps/{app_id}/users/{user_id}/sessions/{session_id}/metamessages/{metamessage_id}
---
openapi: put /v1/apps/{app_id}/users/{user_id}/metamessages/{metamessage_id}
---

4
docs/api-reference/endpoint/sessions/get-session.mdx
View File

@ -1,3 +1,3 @@
---
openapi: get /v1/apps/{app_id}/users/{user_id}/sessions/{session_id}
---
openapi: get /v1/apps/{app_id}/users/{user_id}/sessions
---

4
docs/api-reference/endpoint/users/get-user.mdx
View File

@ -1,3 +1,3 @@
---
openapi: get /v1/apps/{app_id}/users/{user_id}
---
openapi: get /v1/apps/{app_id}/users
---

4
docs/contributing/guidelines.mdx
View File

@ -14,4 +14,6 @@ a project manager. Feel free to join us in our
[discord](http://discord.gg/plasticlabs) to discuss your changes or get help.
Your changes will undergo a period of testing and discussion before finally
being entered into the `main` branch and being staged for release
being entered into the `main` branch and being staged for release. For more
details, check out our [contributing](https://github.com/plastic-labs/honcho/blob/main/CONTRIBUTING.md)
document.

14
docs/contributing/self-hosting.mdx
View File

@ -18,7 +18,7 @@ The minimum poetry version is `0.4.9`
### Setup
Once the dependencies are installed on the system run the following steps to get
the local project setup.
the local project setup.
1. Clone the repository
@ -49,14 +49,14 @@ source honcho/.venv/bin/activate
3. Set up a database
Honcho utilized [Postgres](https://www.postgresql.org/) for its database with
Honcho utilized [Postgres](https://www.postgresql.org/) for its database with
pgvector. An easy way to get started with a postgresdb is to create a project
with [Supabase](https://supabase.com/)
A `docker-compose` template is also available with a database configuration
available.
available.
4. Edit the environment variables.
4. Edit the environment variables.
Honcho uses a `.env` file for managing runtime environment variables. A
`.env.template` file is included for convenience. Several of the configurations
@ -75,10 +75,10 @@ ANTHROPIC_API_KEY= # API Key for Anthropic used for the deriver and dialectic AP
The template has the additional functionality disabled by default. To ensure
that they are disabled you can verify the following environment variables are
set to false.
set to false.
```env
USE_AUTH_SERVICE=false
USE_AUTH=false
SENTRY_ENABLED=false
```
@ -99,7 +99,7 @@ necessary tables for Honcho to operate.
As mentioned earlier a `docker-compose` template is included for running Honcho.
As an alternative to running Honcho locally it can also be run with the compose
template.
template.
Copy the template and update the appropriate environment variables before
launching the service.

123
docs/getting-started/architecture.mdx
View File

@ -4,22 +4,24 @@ description: 'Learn the core primitives and the structure of Honcho'
icon: 'building'
---
Honcho is a user context management system for AI powered applications. It is
inspired by but not a 1:1 mapping of the OpenAI Assistants API. While we are
building many similar primitives, we are going about it differently.
Honcho is a user context management system for AI powered applications.
The storage concepts are inspired by, but not a 1:1 mapping of, the OpenAI
Assistants API. The insights concepts are inspired by cognitive science,
philosophy, and machine learning.
Honcho is open source. We believe trust and transparency are vital for
developing AI technology. We're also focused on using and supporting existing
tools rather than developing from scratch.
One of the main objectives of the Honcho project is to promote community
exploration of *user models*. Language models are highly capable of modeling
human psychology. By building a data management framework that is user-centric,
we aim to address not only practical application development issues (like
scaling, statefulness, etc.) but also kickstart exploration of the design space
of what's possible in terms of building user models. You can read more about
Honcho's origin, inspiration and philosophy on our
[blog](https://blog.plasticlabs.ai).
We focus on flexible, user-centric storage primitives to promote community
exploration of novel memory frameworks and the usage of the
[Dialectic API](https://blog.plasticlabs.ai/blog/Introducing-Honcho's-Dialectic-API)
to support them. Language models are highly capable of modeling human psychology.
By building a data management framework that is user-centric, we aim to address
not only practical application development issues (like scaling, statefulness,
etc.) but also kickstart exploration of the design space of what's possible
given access to rich user models. You can read more about Honcho's origin,
inspiration and philosophy on our [blog](https://blog.plasticlabs.ai).
## Core Primitives
@ -29,11 +31,92 @@ Using Honcho has the following flow:
3. Create a `Session` for a `User`.
4. Create a `Collection` for a `User`
5. Add `Message`s to a `User`'s `Session`.
6. Add `Metamessage`s to `User`'s `Message`s
6. Add `Metamessage`s to a `User` (optional links to `Session`, `Message`)
7. Add `Document`s to a `User`'s `Collection`
![DB Schema](../images/db_schema.png)
```mermaid
erDiagram
App ||--o{ User : "has"
User ||--o{ Session : "has"
User ||--o{ Collection : "has"
User ||--o{ Metamessage : "has"
Session ||--o{ Message : "contains"
Session ||--o{ Metamessage : "has"
Message ||--o{ Metamessage : "has"
Collection ||--o{ Document : "contains"
App {
BigInteger id PK
string public_id
string name
datetime created_at
jsonb h_metadata "metadata"
}
User {
BigInteger id PK
string public_id
string name
jsonb h_metadata "metadata"
datetime created_at
string app_id FK
}
Session {
BigInteger id PK
string public_id
boolean is_active
jsonb h_metadata "metadata"
datetime created_at
string user_id FK
}
Message {
BigInteger id PK
string public_id
string session_id FK
boolean is_user
string content
jsonb h_metadata "metadata"
datetime created_at
}
Metamessage {
BigInteger id PK
string public_id
string metamessage_type
string content
string user_id FK
string session_id FK "nullable"
string message_id FK "nullable"
datetime created_at
jsonb h_metadata "metadata"
}
Collection {
BigInteger id PK
string public_id
string name
datetime created_at
jsonb h_metadata "metadata"
string user_id FK
}
Document {
BigInteger id PK
string public_id
jsonb h_metadata "metadata"
string content
vector embedding "1536"
datetime created_at
string collection_id FK
}
```
### Apps
An `App` is the highest-level primitive in Honcho. It is the scope that all of your `Users` are bound to.
### Users
@ -54,9 +137,15 @@ This is pretty straightforward.
### Metamessages
Plenty of applications have intermediate steps between `User` input and the
response that gets sent. The `Metamessage` object allows you to store those
intermediate steps and link them to the messages they were derived from.
Success in LLM applications is dependent on elegant context management, so we
provide a `Metamessage` object for flexible context storage and construction. Each
`Metamessage` is tied to a `User` object via the required `user_id` argument. Keeping
this separate from the core user-assistant message history ensures the
insights service running ambiently is doing so on authentic ground truth
We've found this particularly useful for storing intermediate inferences,
constructing very specific chat histories, and more. Metamessages can optionally be
attached to sessions and/or messages, so constructing historical context for inference is
as easy as possible.
### Collections
@ -80,4 +169,4 @@ developers implement their own solutions that without a doubt increase overhead
and maintenance. Our goal with Honcho is to provide a simple and flexible
storage framework accompanied by a smooth developer experience to ease pains
building the cumbersome parts of LLM applications. We hope this will allow
developers more freedom to explore exciting, yet-to-be-discovered areas!
developers more freedom to explore exciting, yet-to-be-discovered areas!

88
docs/guides/dialectic-endpoint.mdx
View File

@ -1,38 +1,25 @@
---
title: "Dialectic Endpoint"
description: "An endpoint for easily reasoning about your users"
description: "An endpoint for reasoning about your users"
icon: "comments"
---
> This guide goes over automatic insights generated by Honcho. An example
> of this being used can be found in [Curation Buddy](https://github.com/vintrocode/curation-buddy)
Honcho will do automatic reasoning for you to derive facts about users and
allow your own agents to use them to reason about the user's needs. There are
two aspects to this:
1. Automatic Fact Derivation
2. Dialectic Endpoint
Honcho by default runs ambient inference on top of the `message` objects you store. Those messages serve as the ground truth upon which facts about the user are derived and stored. The **Dialectic Endpoint** is the natural language interface through which insights are synthesized from those facts. We believe [intellectual respect](https://blog.plasticlabs.ai/extrusions/Extrusion-02.24) for LLMs is paramount in building effective AI agents/apps. It follows that the LLM should know better than any human what would aid them in their generation task. Thus, the Dialectic endpoint exists for flexible agent-to-agent communication.
## Automatic Fact Derivation
When you are saving conversations in sessions and messages via Honcho an automatic callback is run that
will reason about the conversations and store facts in a `collection` named **Honcho**. This is a reserved `collection`
specifically for the backend Honcho agent to interact with.
These facts are derived asynchonously and automatically as your users interact with your agents.
On every message written to a session, an automatic callback is run that will reason about the conversation and store facts in a `collection` named `honcho`. This is a reserved `collection` specifically for the backend Honcho agent to interact with.
## Dialectic Endpoint
You can make use the automatically derived facts in the `Honcho` collection directly by querying the documents stored in it,
but an alternative is to use the *Dialectic Endpoint`. What this is, is an endpoint that allows you to talk to an agent that
can automatically take the collection into their context and reason about the users with you.
You can query the automatically derived facts in the `honcho` collection directly, or you can offload this task to our agent and use the Dialectic endpoint. This endpoint allows you to define logic enabling your agent to talk to our agent that automatically retrieves and synthesizes facts from the collection.
This chat interface is exposed via the `Sessions` object.
This chat interface is exposed via the `chat` endpoint. It accepts a string or a list of strings. Below is some example code on how this works.
Belows is some example code on how this works.
## Prerequisites
```python
<CodeGroup>
```python Python
from honcho import Honcho
honcho = Honcho()
@ -46,11 +33,64 @@ user = honcho.apps.users.get_or_create(app_id=app.id, name="demo-user")
# create a new session
session = honcho.apps.users.session.create(app_id=app.id, user_id=user.id)
# Talk to the dialectic agent to reason about their needs
answer = honcho.apps.users.session.chat(app_id=app.id, user_id=user.id, session_id=session.id, query="What is the user's favorite way of completing the task")
# (assuming some messages have been written to Honcho for the deriver to use)
```
```javascript NodeJS
import Honcho from 'honcho-ai';
const honcho = new Honcho(); // defaults to demo server
// Create or get an existing App
const app = await honcho.apps.getOrCreate('demo-app');
// create or get user
const user = await honcho.apps.users.getOrCreate(app.id, 'demo-user');
// create a new session (need to send empty body because it's a POST request)
const session = await honcho.apps.users.sessions.create(app.id, user.id, {});
// (assuming some messages have been written to Honcho for the deriver to use)
```
</CodeGroup>
## Static Dialectic Call
<CodeGroup>
```python Python
query = "What is the user's favorite way of completing the task?"
answer = honcho.apps.users.session.chat(app_id=app.id, user_id=user.id, session_id=session.id, queries=query)
```
```javascript NodeJS
const query = 'What is the user's favorite way of completing the task?'
const DialecticResponse = await honcho.apps.users.sessions.chat(app.id, user.id, session.id, {
queries: query,
});
```
</CodeGroup>
## Streaming Dialectic Call
<CodeGroup>
```python Python
with honcho.apps.users.sessions.with_streaming_response.stream(
app_id=app.id,
user_id=user.id,
session_id=session.id,
queries="What do we know about the user",
) as response:
print(response)
for line in response.iter_text():
print(line)
time.sleep(0.025)
```
```javascript NodeJS
```
</CodeGroup>
We've designed the Dialectic endpoint to be infinitely flexible. We wrote an incomplete list of ideas on how to use it on our blog [here](https://blog.plasticlabs.ai/blog/Introducing-Honcho's-Dialectic-API#how-it-works).

318
docs/guides/discord.mdx
View File

@ -5,7 +5,7 @@ description: "Discord is a powerful chat application that handles many UI compli
sidebarTitle: 'Discord'
---
> Example code is available on [GitHub](https://github.com/plastic-labs/honcho-python/blob/main/examples/discord/roast-bot/main.py)
> Example code is available on [GitHub](https://github.com/plastic-labs/discord-python-starter)
Any application interface that defines logic based on events and supports
special commands can work easily with Honcho. Here's how to use Honcho with
@ -19,56 +19,83 @@ Most Discord bots have async functions that listen for specific events, the most
```python
@bot.event
async def on_message(message):
"""Event that is run when a message is sent in a channel or DM that the bot has access to"""
global last_message_id
if message.author == bot.user:
# ensure the bot does not reply to itself
return
user_id = f"discord_{str(message.author.id)}"
user = honcho.apps.users.get_or_create(name=user_id, app_id=app.id)
is_dm = isinstance(message.channel, discord.DMChannel)
is_reply_to_bot = (
message.reference and message.reference.resolved.author == bot.user
)
is_mention = bot.user.mentioned_in(message)
# Get the session associated with the user and location
location_id = str(message.channel.id) # Get the channel id for the message
if is_dm or is_reply_to_bot or is_mention:
# Remove the bot's mention from the message content if present
input = message.content.replace(f"<@{bot.user.id}>", "").strip()
sessions = [
session
for session in honcho.apps.users.sessions.list(
user_id=user.id, app_id=app.id, is_active=True, location_id=location_id
# If the message is empty after removing the mention, ignore it
if not input:
return
# Get a user object for the message author
user_id = f"discord_{str(message.author.id)}"
user = honcho.apps.users.get_or_create(name=user_id, app_id=app.id)
# Use the channel ID as the location_id (for DMs, this will be unique to the user)
location_id = str(message.channel.id)
# Get or create a session for this user and location
session, _ = get_session(user.id, location_id, create=True)
# Get messages
history_iter = honcho.apps.users.sessions.messages.list(
app_id=app.id, session_id=session.id, user_id=user.id
)
]
history = list(msg for msg in history_iter)
if len(sessions) > 0:
session = sessions[0]
else:
session = honcho.apps.users.sessions.create(user_id=user.id, app_id=app.id, location_id=location_id)
# Add user message to session
user_msg = honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=input,
is_user=True,
)
last_message_id = user_msg.id
history = [
message
for message in honcho.apps.users.sessions.messages.list(session_id=session.id, app_id=app.id, user_id=user.id)
]
chat_history = messages_to_langchain(history)
async with message.channel.typing():
response = llm(input, history)
inp = message.content
honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=input,
is_user=True,
)
if len(response) > 1500:
# Split response into chunks at newlines, keeping under 1500 chars
chunks = []
current_chunk = ""
for line in response.splitlines(keepends=True):
if len(current_chunk) + len(line) > 1500:
chunks.append(current_chunk)
current_chunk = line
else:
current_chunk += line
if current_chunk:
chunks.append(current_chunk)
for chunk in chunks:
await message.channel.send(chunk)
else:
await message.channel.send(response)
async with message.channel.typing():
response = await chain.ainvoke({"chat_history": chat_history, "input": inp})
await message.channel.send(response)
honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=response,
is_user=False,
)
# Add bot message to session
honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=response,
is_user=False,
)
```
Let's break down what each chunk of code is doing...
Let's break down what this code is doing...
```python
@bot.event
async def on_message(message):
@ -79,67 +106,114 @@ async def on_message(message):
This is how you define an event function in `py-cord` that listens for messages and checks that the bot doesn't respond to itself.
```python
is_dm = isinstance(message.channel, discord.DMChannel)
is_reply_to_bot = (
message.reference and message.reference.resolved.author == bot.user
)
is_mention = bot.user.mentioned_in(message)
```
These lines check what kind of message is being sent in Discord, which is a useful condition to check before entering the reply logic. The code inside that if-statement is commented quite well, so we'll just go over the relevant Honcho parts.
```python
# Get a user object for the message author
user_id = f"discord_{str(message.author.id)}"
user = honcho.apps.users.get_or_create(name=user_id, app_id=app.id)
```
Here we're getting or creating a user for an app that's been defined at the top of the file.
```python
# Use the channel ID as the location_id (for DMs, this will be unique to the user)
location_id = str(message.channel.id)
# Get or create a session for this user and location
session, _ = get_session(user.id, location_id, create=True)
```
Honcho accepts a `location_id` argument to help separate out locations messages were sent (which is convenient for Discord channels).
Here we're using the discord channel ID as a unique `location_id` to attach as metadata to the session. Then we have a nice [helper function](https://github.com/plastic-labs/discord-python-starter/blob/main/src/bot.py#L85) to take care of some of the session querying logic--we'll dive into that shortly.
```python
sessions = [
session
for session in honcho.apps.users.sessions.list(
user_id=user.id, app_id=app.id, is_active=True, location_id=location_id
# Get messages
history_iter = honcho.apps.users.sessions.messages.list(
app_id=app.id, session_id=session.id, user_id=user.id
)
history = list(msg for msg in history_iter)
```
When you call the `list` method, it returns an iterable which you can quickly loop over to create a list of `Message` objects. Then, we make the call to the LLM using another neat [helper function](https://github.com/plastic-labs/discord-python-starter/blob/main/src/bot.py#L52) that we will cover.
## Helper functions
The first helper function we create is called `get_session`. This simplifies a lot of our session-querying logic.
```python
def get_session(user_id, location_id, create=False):
# Get an existing session for the user and location or optionally create a new one if none exists.
# Returns a tuple of (session, is_new) where is_new indicates if a new session was created.
# Query for *active* sessions with both user_id and location_id
sessions_iter = honcho.apps.users.sessions.list(
app_id=app.id, user_id=user_id, reverse=True, is_active=True
)
sessions = list(session for session in sessions_iter)
# Find the right session
for session in sessions:
if session.metadata.get("location_id") == location_id:
return session, False
# If no session is found and create is True, create a new one
if create:
print("No active session found, creating new one")
return honcho.apps.users.sessions.create(
user_id=user_id,
app_id=app.id,
metadata={"location_id": location_id},
), True
return None, False
```
You can see the `list` method on the sessions object similarly returns an iterable. This is a common pattern in Honcho--use list comprehension to create your new python list. Then loop through those session objects to find the appropriate `location_id` in the metadata, and if none are found then create a new session. You'll also notice we list messages in `reverse=True` order--this means you will get the most recent ones first. We also support native filtering by active sessions.
The next helper function we create is called `llm`. This simplifies constructing the chat message object we're going to send to the inference provider.
```python
def llm(prompt, previous_chats=None):
messages = []
# Add system message with documentation context
messages.append(
{
"role": "system",
"content": f"You are a helpful assistant."
}
)
if previous_chats:
messages.extend(
[
{"role": "user" if msg.is_user else "assistant", "content": msg.content}
for msg in previous_chats
]
)
]
if len(sessions) > 0:
session = sessions[0]
else:
session = honcho.apps.users.sessions.create(user_id=user.id, app_id=app.id, location_id=location_id)
messages.append({"role": "user", "content": prompt})
try:
completion = openai.chat.completions.create(
model=MODEL_NAME,
messages=messages,
)
return completion.choices[0].message.content
except Exception as e:
print(e)
return f"Error: {e}"
```
Here we're querying honcho for the user's sessions based on the location (channel) they're in. This will get all the sessions, so the if statement just pops the most recent one (if there are many) or creates a new one if none exist.
```python
history = [
message
for message in honcho.apps.users.sessions.messages.list(session_id=session.id, app_id=app.id, user_id=user.id)
]
chat_history = messages_to_langchain(history)
# Add user message to session
input = message.content
honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=input,
is_user=True,
)
async with message.channel.typing():
response = await chain.ainvoke({"chat_history": chat_history, "input": inp})
await message.channel.send(response)
# Add bot message to session
honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=response,
is_user=False,
)
```
This chunk is all about constructing the object to send to an LLM API. We get
the messages from a `session` and construct a `chat_history` object with a
quick utility function (more on that in the [Langchain](./langchain) guide).
Then, we access the user message via `message.content` and add it to Honcho.
The `async with` method allows the bot to show that it's "typing" while waiting
for an LLM response and then uses `message.channel.send` to respond to the
user. We can then add that AI response to Honcho with the same
`session.create_message` method, this time specifying that this message did not
come from a user with `is_user=False`.
Note that `messages` is a list of dictionaries that are individually defined with key-value pairs for roles and content. We again use list comprehension to unpack historical message objects into the list that we send to the chat completions method. Honcho `Message` objects store role and content natively to make this context construction as simple as possible. If you're interested in learning more about native Honcho objects, you can check out the [`models.py`](https://github.com/plastic-labs/honcho/blob/main/src/models.py) file.
## Slash Commands
@ -147,38 +221,40 @@ Discord bots also offer slash command functionality. We can use Honcho to do
interesting things via slash commands. Here's a simple example:
```python
@bot.slash_command(name = "restart", description = "Restart the Conversation")
@bot.slash_command(
name="restart",
description="Reset all of your messaging history with Honcho in this channel.",
)
async def restart(ctx):
user_id=f"discord_{str(ctx.author.id)}"
user = honcho.apps.users.get_or_create(name=user_id, app_id=app.id)
location_id=str(ctx.channel_id)
sessions = [
session
for session in honcho.apps.users.sessions.list(
user_id=user.id, app_id=app.id, is_active=True, location_id=location_id
)
]
if len(sessions) > 0:
honcho.apps.users.sessions.delete(app_id=app.id, user_id=user.id, session_id=sessions[0].id)
print(f"restarting conversation for {ctx.author.name}")
async with ctx.typing():
user_name = f"discord_{str(ctx.author.id)}"
user = honcho.apps.users.get_or_create(name=user_name, app_id=app.id)
location_id = str(ctx.channel_id)
# Get existing session
session, _ = get_session(user.id, location_id, create=False)
if session:
# Delete the session
honcho.apps.users.sessions.delete(
app_id=app.id, user_id=user.id, session_id=session.id
)
msg = "The conversation has been restarted."
msg = "Great! The conversation has been restarted. What would you like to talk about?"
await ctx.respond(msg)
```
This slash command restarts a conversation with a bot. In that case, we want to remove that session from storage. You can see we follow the same steps to access the user metadata via commands from the application interface:
```python
user_id=f"discord_{str(ctx.author.id)}"
user = honcho.apps.users.get_or_create(name=user_id, app_id=app.id)
location_id=str(ctx.channel_id)
```
Then we can retrieve and delete the session associated with that metadata:
```python
sessions = [
session
for session in honcho.apps.users.sessions.list(
user_id=user.id, app_id=app.id, is_active=True, location_id=location_id
)
]
if len(sessions) > 0:
honcho.apps.users.sessions.delete(app_id=app.id, user_id=user.id, session_id=sessions[0].id)
```
This slash command restarts a conversation with a bot. In Honcho, the `delete` method marks a session's `is_active` field to `False`.
## Recap
How you use Honcho is tightly coupled with the client you're building in. Here, Discord serves as an example of an interactive chat interface. We're just scratching the surface of things you can do with Honcho, but we learned some key patterns:
- how to register users
- how to work with iterables when listing sessions, messages
- how to attach metadata to Honcho objects (like `location_id` on sessions)
- how to sort and filter when calling list methods
You are well on your way to becoming a context construction master! Stay tuned for more in-depth examples. If you want a challenge, try deciphering how we construct context in one of our apps, [Bloom](https://github.com/plastic-labs/tutor-gpt/blob/main/app/Chat.tsx).

62
docs/guides/honcho-mcp.mdx Normal file
View File

@ -0,0 +1,62 @@
---
title: "Honcho MCP"
icon: 'star-of-life'
description: "Use Honcho in Claude Desktop"
sidebarTitle: 'Claude Desktop'
---
You can let Claude use Honcho to manage its own memory in the native desktop app by using the Honcho MCP integration! Follow these steps:
1. Clone the `honcho-mcp` repo:
```
git clone git@github.com:plastic-labs/honcho-mcp.git
```
2. Navigate into the `honcho-mcp` folder.
```
cd honcho-mcp
```
3. Sync the virtual environment. This package uses [uv](https://docs.astral.sh/uv/), [install](https://docs.astral.sh/uv/#installation) if you haven't.
```
uv sync
```
4. In Claude Desktop, go to the *top left Mac Toolbar* Settings > Developer and click "Edit Config"
5. Add the following (and update paths!):
```
{
"mcpServers": {
"Honcho": {
"command": "/path/to/uv",
"args": [
"run",
"--with",
"mcp[cli]",
"mcp",
"run",
"/path/to/honcho-mcp/main.py"
]
}
}
}
```
<Warning>You probably will need to put the full path to the uv executable in the command field. You can get this by running `which uv` on MacOS/Linux or `where uv` on Windows.</Warning>
6. Restart the Claude Desktop app. Upon relaunch it should start Honcho and the tools should be available!
Just note that, by default, the MCP server is set up to use the Honcho Demo server, which only persists data for 7 days. If you're using the hosted version of Honcho, copy the `.env.template` to a proper `.env` file and update the URL and API key variables accordingly.
## Project Instructions
Finally, Claude needs instructions on how to use Honcho. The Desktop app doesn't allow you to add system prompts directly, but you can create a project and paste these [instructions](https://github.com/plastic-labs/honcho-mcp/blob/main/instructions.txt) into the "Project Instructions" field.
<Note>Be sure to update the \<app_name\> and \<user_name\> variables in the instructions.txt file.</Note>
Claude should then query for insights before responding and write your messages to storage! If you come up with more creative ways to get Claude to manage its own memory with Honcho, feel free to [let us know](https://discord.gg/plasticlabs) or make a PR on this [repo](https://github.com/plastic-labs/honcho-mcp/tree/main)!

94
docs/guides/langchain.mdx
View File

@ -1,94 +0,0 @@
---
title: 'LangChain Integration 🦜⛓️'
sidebarTitle: 'LangChain'
description: 'Using Honcho with LangChain with drop-in primitives'
icon: 'bird'
---
You can use Honcho to manage user context around LLM frameworks like LangChain. First, import the appropriate packages:
```python
from honcho import Honcho
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import AIMessage, HumanMessage
from dotenv import load_dotenv # for loading in LLM API keys
load_dotenv() # assumes you have a .env file with OPENAI_API_KEY defined
```
Next let's instantiate our Honcho client:
```python
honcho = Honcho(environment="demo")
app_name = "LangChain App"
app = honcho.apps.get_or_create(name=app_name) # create or get app
user_name = str(uuid4())
user = honcho.apps.users.get_or_create(app_id=app.id, name=user_name) # create or get user
```
Then we can define our chain using the LangChain Expression Language ([LCEL](https://python.langchain.com/docs/expression_language/why)):
```python
prompt = ChatPromptTemplate.from_messages([
("system", "You are a helpful assistant."),
MessagesPlaceholder(variable_name="chat_history"),
("user", "{input}")
])
model = ChatOpenAI(model="gpt-3.5-turbo")
output_parser = StrOutputParser()
chain = prompt | model | output_parser
```
Honcho returns lists of `Message` objects when queried using a built-in method like `get_messages()`, so a quick utility function is needed to change the list format to message objects LangChain expects:
```python
def messages_to_langchain(messages: List):
new_messages = []
for message in messages:
if message.is_user:
new_messages.append(HumanMessage(content=message.content))
else:
new_messages.append(AIMessage(content=message.content))
return new_messages
```
This method is importable with the following statement
```python
from honcho.lib.ext.langchain import messages_to_langchain
```
Now we can structure Honcho calls around our LLM inference:
```python
sessions = [
session for session in
honcho.apps.users.sessions.list(
app_id=app.id,
user_id=user.id,
location_id=location_id
)
] # args come from application logic
session = sessions[0] # most recent session for user
history = [
message for message in
honcho.apps.users.sessions.messages.list(
app_id=app.id,
user_id=user.id,
session_id=session.id
)
]
chat_history = messages_to_langchain(history) # convert messages for LangChain
inp = "Here's a user message!"
honcho.apps.users.sessions.messages.create(app_id=app.id, user_id=user.id, is_user=True, content=inp)
response = await chain.ainvoke({"chat_history": chat_history, "input": inp})
honcho.apps.users.sessions.messages.create(app_id=app.id, user_id=user.id, is_user=False, content=response)
```
Here we query messages from a user's session using Honcho and construct a chat history object to send to the LLM alongside our immediate user input. Once the LLM has responded, we can add that to Honcho!

240
docs/guides/streaming-response.mdx Normal file
View File

@ -0,0 +1,240 @@
---
title: "Streaming Responses"
description: "Using streaming responses with Honcho SDKs"
icon: "wave-sine"
---
# Streaming Responses with Honcho
When working with AI-generated content, streaming the response as it's generated can significantly improve the user experience. Honcho provides streaming functionality in its SDKs that allows your application to display content as it's being generated, rather than waiting for the complete response.
## When to Use Streaming
Streaming is particularly useful for:
- Real-time chat interfaces
- Long-form content generation
- Applications where perceived speed is important
- Interactive agent experiences
- Reducing time-to-first-word in user interactions
## Streaming with the Dialectic Endpoint
One of the primary use cases for streaming in Honcho is with the Dialectic endpoint. This allows you to stream the AI's reasoning about a user in real-time.
### Prerequisites
<CodeGroup>
```python Python
from honcho import Honcho
honcho = Honcho()
# Create or get an existing App
app = honcho.apps.get_or_create(name="demo-app")
# Create or get user
user = honcho.apps.users.get_or_create(app_id=app.id, name="demo-user")
# Create a new session
session = honcho.apps.users.sessions.create(app_id=app.id, user_id=user.id)
# Store some messages for context (optional)
honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content="Hello, I'm testing the streaming functionality",
is_user=True
)
```
```javascript NodeJS
import Honcho from 'honcho-ai';
const honcho = new Honcho();
// Create or get an existing App
const app = await honcho.apps.getOrCreate('demo-app');
// Create or get user
const user = await honcho.apps.users.getOrCreate(app.id, 'demo-user');
// Create a new session
const session = await honcho.apps.users.sessions.create(app.id, user.id, {});
// Store some messages for context (optional)
await honcho.apps.users.sessions.messages.create(app.id, user.id, session.id, {
content: "Hello, I'm testing the streaming functionality",
is_user: true
});
```
</CodeGroup>
## Streaming from the Dialectic Endpoint
<CodeGroup>
```python Python
import time
# Basic streaming example
with honcho.apps.users.sessions.with_streaming_response.stream(
app_id=app.id,
user_id=user.id,
session_id=session.id,
queries="What can you tell me about this user?",
) as response:
for chunk in response.iter_text():
print(chunk, end="", flush=True) # Print each chunk as it arrives
time.sleep(0.01) # Optional delay for demonstration
```
```javascript NodeJS
// Basic streaming example
const stream = await honcho.apps.users.sessions.chat(app.id, user.id, session.id, {
queries: "What can you tell me about this user?",
stream: true
});
// Process the stream
for await (const chunk of stream) {
process.stdout.write(chunk); // Write to console without newlines
}
```
</CodeGroup>
## Working with Streaming Data
When working with streaming responses, consider these patterns:
1. **Progressive Rendering** - Update your UI as chunks arrive instead of waiting for the full response
2. **Buffered Processing** - Accumulate chunks until a logical break (like a sentence or paragraph)
3. **Token Counting** - Monitor token usage in real-time for applications with token limits
4. **Error Handling** - Implement appropriate error handling for interrupted streams
## Example: Restaurant Recommendation Chat
<CodeGroup>
```python Python
import asyncio
from honcho import Honcho
async def restaurant_recommendation_chat():
honcho = Honcho()
app = await honcho.apps.get_or_create(name="food-app")
user = await honcho.apps.users.get_or_create(app_id=app.id, name="food-lover")
session = await honcho.apps.users.sessions.create(app_id=app.id, user_id=user.id)
# Store multiple user messages about food preferences
user_messages = [
"I absolutely love spicy Thai food, especially curries with coconut milk.",
"Italian cuisine is another favorite - fresh pasta and wood-fired pizza are my weakness!",
"I try to eat vegetarian most of the time, but occasionally enjoy seafood.",
"I can't handle overly sweet desserts, but love something with dark chocolate."
]
# Store the user's messages in the session
for message in user_messages:
await honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=message,
is_user=True
)
print(f"User: {message}")
# Ask for restaurant recommendations based on preferences
print("\nRequesting restaurant recommendations...")
print("Assistant: ", end="", flush=True)
full_response = ""
# Stream the response
with honcho.apps.users.sessions.with_streaming_response.stream(
app_id=app.id,
user_id=user.id,
session_id=session.id,
queries="Based on this user's food preferences, recommend 3 restaurants they might enjoy in the Lower East Side."
) as response:
for chunk in response.iter_text():
print(chunk, end="", flush=True)
full_response += chunk
await asyncio.sleep(0.01)
# Store the assistant's complete response
await honcho.apps.users.sessions.messages.create(
app_id=app.id,
user_id=user.id,
session_id=session.id,
content=full_response,
is_user=False
)
# Run the async function
if __name__ == "__main__":
asyncio.run(restaurant_recommendation_chat())
```
```javascript NodeJS
import Honcho from 'honcho-ai';
async function restaurantRecommendationChat() {
const honcho = new Honcho();
const app = await honcho.apps.getOrCreate('food-app');
const user = await honcho.apps.users.getOrCreate(app.id, 'food-lover');
const session = await honcho.apps.users.sessions.create(app.id, user.id, {});
// Store multiple user messages about food preferences
const userMessages = [
"I absolutely love spicy Thai food, especially curries with coconut milk.",
"Italian cuisine is another favorite - fresh pasta and wood-fired pizza are my weakness!",
"I try to eat vegetarian most of the time, but occasionally enjoy seafood.",
"I can't handle overly sweet desserts, but love something with dark chocolate."
];
// Store the user's messages in the session
for (const message of userMessages) {
await honcho.apps.users.sessions.messages.create(app.id, user.id, session.id, {
content: message,
is_user: true
});
console.log(`User: ${message}`);
}
// Ask for restaurant recommendations based on preferences
console.log("\nRequesting restaurant recommendations...");
process.stdout.write("Assistant: ");
let fullResponse = "";
// Stream the response
const stream = await honcho.apps.users.sessions.chat(app.id, user.id, session.id, {
queries: "Based on this user's food preferences, recommend 3 restaurants they might enjoy in the Lower East Side.",
stream: true
});
for await (const chunk of stream) {
process.stdout.write(chunk);
fullResponse += chunk;
}
// Store the assistant's complete response
await honcho.apps.users.sessions.messages.create(app.id, user.id, session.id, {
content: fullResponse,
is_user: false
});
}
restaurantRecommendationChat().catch(console.error);
```
</CodeGroup>
## Performance Considerations
When implementing streaming:
- Consider connection stability for mobile or unreliable networks
- Implement appropriate timeouts for stream operations
- Be mindful of memory usage when accumulating large responses
- Use appropriate error handling for network interruptions
Streaming responses provide a more interactive and engaging user experience. By implementing streaming in your Honcho applications, you can create more responsive AI-powered features that feel natural and immediate to your users.

21
docs/mint.json
View File

@ -61,19 +61,15 @@
},
{
"group": "Getting Started",
"pages": ["guides/overview"]
"pages": ["guides/overview", "guides/streaming-response"]
},
{
"group": "Application Interfaces",
"pages": ["guides/discord"]
},
{
"group": "Integrations",
"pages": ["guides/langchain"]
"pages": ["guides/discord", "guides/honcho-mcp"]
},
{
"group": "Personal Memory",
"pages": ["guides/simple-memory", "guides/dialectic-endpoint"]
"pages": ["guides/dialectic-endpoint"]
},
{
"group": "API Documentation",
@ -83,6 +79,7 @@
"group": "apps",
"pages": [
"api-reference/endpoint/apps/get-app",
"api-reference/endpoint/apps/get-all-apps",
"api-reference/endpoint/apps/update-app",
"api-reference/endpoint/apps/get-app-by-name",
"api-reference/endpoint/apps/create-app",
@ -109,7 +106,6 @@
"api-reference/endpoint/sessions/update-session",
"api-reference/endpoint/sessions/delete-session",
"api-reference/endpoint/sessions/chat",
"api-reference/endpoint/sessions/chat-stream",
"api-reference/endpoint/sessions/clone-session"
]
},
@ -129,8 +125,7 @@
"api-reference/endpoint/metamessages/create-metamessage",
"api-reference/endpoint/metamessages/get-metamessages",
"api-reference/endpoint/metamessages/get-metamessage",
"api-reference/endpoint/metamessages/update-metamessage",
"api-reference/endpoint/metamessages/get-metamessages-by-user"
"api-reference/endpoint/metamessages/update-metamessage"
]
},
{
@ -139,7 +134,7 @@
"api-reference/endpoint/collections/get-collections",
"api-reference/endpoint/collections/create-collection",
"api-reference/endpoint/collections/get-collection-by-name",
"api-reference/endpoint/collections/get-collection-by-id",
"api-reference/endpoint/collections/get-collection",
"api-reference/endpoint/collections/update-collection",
"api-reference/endpoint/collections/delete-collection"
]
@ -154,6 +149,10 @@
"api-reference/endpoint/documents/delete-document",
"api-reference/endpoint/documents/query-documents"
]
},
{
"group": "keys",
"pages": ["api-reference/endpoint/keys/create-key"]
}
],
"footerSocials": {

2624
docs/openapi.json
View File

File diff suppressed because it is too large Load Diff

19701
docs/package-lock.json generated
View File

File diff suppressed because it is too large Load Diff

1
docs/package.json
View File

@ -12,6 +12,7 @@
"license": "ISC",
"dependencies": {
"@mintlify/scraping": "^3.0.92",
"honcho-ai": "^0.0.11",
"mintlify": "^4.0.245"
}
}

2
migrations/env.py
View File

@ -95,7 +95,7 @@ def run_migrations_online() -> None:
connectable = engine_from_config(
configuration,
prefix="sqlalchemy.",
echo=True,
echo=False,
poolclass=pool.NullPool,
)

161
migrations/versions/b765d82110bd_change_metamessages_to_user_level_with_.py Normal file
View File

@ -0,0 +1,161 @@
"""Change Metamessages to user level with optional message and session link
Revision ID: b765d82110bd
Revises: c3828084f472
Create Date: 2025-04-03 15:32:16.733312
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.sql import text
from sqlalchemy.dialects import postgresql
# revision identifiers, used by Alembic.
revision: str = "b765d82110bd"
down_revision: Union[str, None] = "c3828084f472"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# 1. Add new columns to metamessages table
op.add_column("metamessages", sa.Column("user_id", sa.TEXT(), nullable=True))
op.add_column("metamessages", sa.Column("session_id", sa.TEXT(), nullable=True))
# 2. Create foreign key constraints for the new columns
op.create_foreign_key(
"fk_metamessages_user_id_users",
"metamessages",
"users",
["user_id"],
["public_id"],
)
op.create_foreign_key(
"fk_metamessages_session_id_sessions",
"metamessages",
"sessions",
["session_id"],
["public_id"],
)
# 3. Make message_id nullable
op.alter_column(
"metamessages", "message_id", existing_type=sa.TEXT(), nullable=True
)
# 4. Create indices for users, sessions, and messages - only if they don't exist
conn = op.get_bind()
inspector = sa.inspect(conn)
# Check and create each index only if it doesn't exist
existing_indices = {
idx["name"]
for schema in [None, "public"]
for tbl in inspector.get_table_names(schema=schema)
for idx in inspector.get_indexes(tbl, schema=schema)
}
# Helper function to create index if not exists
def create_index_if_not_exists(index_name, table_name, columns, **kwargs):
if index_name not in existing_indices:
op.create_index(index_name, table_name, columns, **kwargs)
# Create all indices with the helper function
create_index_if_not_exists("idx_users_app_lookup", "users", ["app_id", "public_id"])
create_index_if_not_exists(
"idx_sessions_user_lookup", "sessions", ["user_id", "public_id"]
)
# For complex indices
create_index_if_not_exists(
"idx_messages_session_lookup",
"messages",
["session_id", "id"],
postgresql_include=["public_id", "is_user", "created_at"],
)
create_index_if_not_exists(
"idx_metamessages_lookup",
"metamessages",
["metamessage_type", text("id DESC")],
postgresql_include=["public_id", "message_id", "created_at"],
)
create_index_if_not_exists(
"idx_metamessages_user_lookup",
"metamessages",
["user_id", "metamessage_type", text("id DESC")],
)
create_index_if_not_exists(
"idx_metamessages_session_lookup",
"metamessages",
["session_id", "metamessage_type", text("id DESC")],
)
create_index_if_not_exists(
"idx_metamessages_message_lookup",
"metamessages",
["message_id", "metamessage_type", text("id DESC")],
)
# 7. Add the check constraint for message_id and session_id relationship
op.create_check_constraint(
"message_requires_session",
"metamessages",
"(message_id IS NULL) OR (session_id IS NOT NULL)",
)
# 8. Update existing data: fill user_id from message's session's user
# This is a complex data migration that requires SQL
op.execute("""
UPDATE metamessages m
SET user_id = u.public_id,
session_id = s.public_id
FROM messages msg
JOIN sessions s ON msg.session_id = s.public_id
JOIN users u ON s.user_id = u.public_id
WHERE m.message_id = msg.public_id
AND m.user_id IS NULL
""")
# 9. Now that data is migrated, make user_id not nullable
op.alter_column("metamessages", "user_id", existing_type=sa.TEXT(), nullable=False)
def downgrade() -> None:
# 1. Remove the check constraint
op.drop_constraint("message_requires_session", "metamessages", type_="check")
# 2. Drop all the new indices
op.drop_index("idx_metamessages_message_lookup", table_name="metamessages")
op.drop_index("idx_metamessages_session_lookup", table_name="metamessages")
op.drop_index("idx_metamessages_user_lookup", table_name="metamessages")
op.drop_index("idx_metamessages_lookup", table_name="metamessages")
op.drop_index("idx_messages_session_lookup", table_name="messages")
op.drop_index("idx_sessions_user_lookup", table_name="sessions")
op.drop_index("idx_users_app_lookup", table_name="users")
# 3. Remove foreign key constraints
op.drop_constraint(
"fk_metamessages_session_id_sessions", "metamessages", type_="foreignkey"
)
op.drop_constraint(
"fk_metamessages_user_id_users", "metamessages", type_="foreignkey"
)
# 4. Make message_id required again and clean up data if needed
op.execute("""
DELETE FROM metamessages WHERE message_id IS NULL
""")
op.alter_column(
"metamessages", "message_id", existing_type=sa.TEXT(), nullable=False
)
# 5. Drop the new columns
op.drop_column("metamessages", "session_id")
op.drop_column("metamessages", "user_id")

9
pyproject.toml
View File

@ -1,6 +1,6 @@
[project]
name = "honcho"
version = "0.0.16"
version = "1.0.0"
description = "Honcho Server"
authors = [
{name = "Plastic Labs", email = "hello@plasticlabs.ai"},
@ -23,6 +23,8 @@ dependencies = [
"nanoid>=2.0.0",
"alembic>=1.14.0",
"langfuse>=2.57.1",
"pyjwt>=2.10.0",
"google-genai>=1.10.0",
]
[tool.uv]
dev-dependencies = [
@ -32,6 +34,7 @@ dev-dependencies = [
"coverage>=7.6.0",
"interrogate>=1.7.0",
"py-spy>=0.3.14",
"ruff>=0.11.2",
]
[tool.ruff.lint]
@ -50,9 +53,9 @@ select = [
# isort
"I",
]
ignore = ["E501"]
ignore = ["E501", "B008"]
[tool.ruff.flake8-bugbear]
[tool.ruff.lint.flake8-bugbear]
extend-immutable-calls = ["fastapi.Depends"]
[tool.pytest.ini_options]

0
scripts/__init__.py Normal file
View File

38
scripts/generate_jwt_secret.py Executable file
View File

@ -0,0 +1,38 @@
#!/usr/bin/env python
"""
Utility script to generate a JWT secret for use in the .env file.
This uses the same logic as the automatically generated version in security.py.
"""
import argparse
import secrets
def generate_jwt_secret():
"""Generate a random JWT secret using the secrets module."""
return secrets.token_hex(32)
def main():
parser = argparse.ArgumentParser(
description="Generate a JWT secret for authentication."
)
parser.add_argument(
"--print-only",
action="store_true",
help="Only print the secret without instructions",
)
args = parser.parse_args()
secret = generate_jwt_secret()
if args.print_only:
print(secret)
else:
print(f"Generated JWT secret: {secret}")
print("\nAdd this to your .env file as:")
print(f"AUTH_JWT_SECRET={secret}")
if __name__ == "__main__":
main()

18
scripts/provision_db.py Normal file
View File

@ -0,0 +1,18 @@
# honcho/scripts/provision_db.py
import os
import sys
# Add the project root to the path
# This assumes the script is run from the scripts directory
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, project_root)
## First import the models to register them with Base
from src import (
models,
) # This registers all models with Base Now you can import from src
from src.db import scaffold_db
if __name__ == "__main__":
scaffold_db()
print("Database created")

542
src/agent.py
View File

@ -1,9 +1,12 @@
import asyncio
import json
import logging
import os
from collections.abc import Iterable
from typing import Any, Optional
import sentry_sdk
from anthropic import Anthropic, MessageStreamManager
from anthropic import MessageStreamManager
from dotenv import load_dotenv
from langfuse.decorators import langfuse_context, observe
from sentry_sdk.ai.monitoring import ai_track
@ -12,6 +15,30 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
from src.db import SessionLocal
from src.deriver.tom import get_tom_inference
from src.deriver.tom.embeddings import CollectionEmbeddingStore
from src.deriver.tom.long_term import get_user_representation_long_term
from src.utils import history, parse_xml_content
from src.utils.model_client import ModelClient, ModelProvider
# Configure logging
logger = logging.getLogger(__name__)
USER_REPRESENTATION_METAMESSAGE_TYPE = "honcho_user_representation"
DEF_DIALECTIC_PROVIDER = ModelProvider.ANTHROPIC
DEF_DIALECTIC_MODEL = "claude-3-7-sonnet-20250219"
DEF_QUERY_GENERATION_PROVIDER = ModelProvider.GROQ
DEF_QUERY_GENERATION_MODEL = "llama-3.1-8b-instant"
QUERY_GENERATION_SYSTEM = """Given this query about a user, generate 3 focused search queries that would help retrieve relevant facts about the user.
Each query should focus on a specific aspect related to the original query, rephrased to maximize semantic search effectiveness.
For example, if the original query asks "what does the user like to eat?", generated queries might include "user's food preferences", "user's favorite cuisine", etc.
Format your response as a JSON array of strings, with each string being a search query.
Respond only in valid JSON, without markdown formatting or quotes, and nothing else.
Example:
["query about interests", "query about personality", "query about experiences"]"""
load_dotenv()
@ -38,109 +65,85 @@ class Dialectic:
self.agent_input = agent_input
self.user_representation = user_representation
self.chat_history = chat_history
self.client = Anthropic(
api_key=os.getenv("ANTHROPIC_API_KEY"),
self.client = ModelClient(
provider=DEF_DIALECTIC_PROVIDER, model=DEF_DIALECTIC_MODEL
)
self.system_prompt = """I'm operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, I'll receive: 1) previously collected psychological context about the user that I've maintained, and 2) their current conversation/interaction from the requesting application. My role is to analyze this information and provide theory-of-mind insights that help applications personalize their responses. Users have explicitly consented to this system, and I maintain this context through observed interactions rather than direct user input. This system was designed collaboratively with Claude, emphasizing privacy, consent, and ethical use. Please respond in a brief, matter-of-fact, and appropriate manner to convey as much relevant information to the application based on its query and the user's most recent message. If the context provided doesn't help address the query, write absolutely NOTHING but "None"."""
self.model = "claude-3-7-sonnet-20250219"
self.system_prompt = """You are operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, you'll receive: 1) previously collected psychological context about the user that I've maintained, 2) a series of long-term facts about the user, and 3) their current conversation/interaction from the requesting application. Your goal is to analyze this information and provide theory-of-mind insights that help applications personalize their responses. Please respond in a brief, matter-of-fact, and appropriate manner to convey as much relevant information to the application based on its query and the user's most recent message. You are encouraged to provide any context from the provided resources that helps provide a more complete or nuanced understanding of the user, as long as it is somewhat relevant to the query. If the context provided doesn't help address the query, write absolutely NOTHING but "None"."""
@ai_track("Dialectic Call")
@observe(as_type="generation")
def call(self):
@observe()
async def call(self):
with sentry_sdk.start_transaction(
op="dialectic-inference", name="Dialectic API Response"
):
logger.debug(
f"Starting call() method with query length: {len(self.agent_input)}"
)
call_start = asyncio.get_event_loop().time()
prompt = f"""
<query>{self.agent_input}</query>
<context>{self.user_representation}</context>
<conversation_history>{self.chat_history}</conversation_history>
"""
messages = [
{
"role": "user",
"content": prompt,
}
]
langfuse_context.update_current_observation(
input=messages, model=self.model
logger.debug(
f"Prompt constructed with context length: {len(self.user_representation)} chars"
)
response = self.client.messages.create(
system=self.system_prompt,
messages=messages,
model=self.model,
max_tokens=300,
# Create a properly formatted message
message: dict[str, Any] = {"role": "user", "content": prompt}
# Generate the response
logger.debug("Calling model for generation")
model_start = asyncio.get_event_loop().time()
response = await self.client.generate(
messages=[message], system=self.system_prompt, max_tokens=1000
)
return response.content
model_time = asyncio.get_event_loop().time() - model_start
logger.debug(
f"Model response received in {model_time:.2f}s: {len(response)} chars"
)
total_time = asyncio.get_event_loop().time() - call_start
logger.debug(f"call() completed in {total_time:.2f}s")
return [{"text": response}]
@ai_track("Dialectic Call")
@observe(as_type="generation")
def stream(self):
@observe()
async def stream(self):
with sentry_sdk.start_transaction(
op="dialectic-inference", name="Dialectic API Response"
):
logger.debug(
f"Starting stream() method with query length: {len(self.agent_input)}"
)
stream_start = asyncio.get_event_loop().time()
prompt = f"""
<query>{self.agent_input}</query>
<context>{self.user_representation}</context>
<conversation_history>{self.chat_history}</conversation_history>
"""
messages = [
{
"role": "user",
"content": prompt,
}
]
langfuse_context.update_current_observation(
input=messages, model=self.model
logger.debug(
f"Prompt constructed with context length: {len(self.user_representation)} chars"
)
return self.client.messages.stream(
model=self.model,
system=self.system_prompt,
messages=messages,
max_tokens=300,
# Create a properly formatted message
message: dict[str, Any] = {"role": "user", "content": prompt}
# Stream the response
logger.debug("Calling model for streaming")
model_start = asyncio.get_event_loop().time()
stream = await self.client.stream(
messages=[message], system=self.system_prompt, max_tokens=1000
)
stream_setup_time = asyncio.get_event_loop().time() - model_start
logger.debug(f"Stream started in {stream_setup_time:.2f}s")
async def chat_history(app_id: str, user_id: str, session_id: str) -> str:
async with SessionLocal() as db:
stmt = await crud.get_messages(db, app_id, user_id, session_id)
results = await db.execute(stmt)
messages = results.scalars()
history = ""
for message in messages:
if message.is_user:
history += f"user:{message.content}\n"
else:
history += f"assistant:{message.content}\n"
return history
async def get_latest_user_representation(
db: AsyncSession, app_id: str, user_id: str
) -> str:
stmt = (
select(models.Metamessage)
.join(models.Message, models.Message.public_id == models.Metamessage.message_id)
.join(models.Session, models.Message.session_id == models.Session.public_id)
.join(models.User, models.User.public_id == models.Session.user_id)
.join(models.App, models.App.public_id == models.User.app_id)
.where(models.App.public_id == app_id)
.where(models.User.public_id == user_id)
.where(models.Metamessage.metamessage_type == "user_representation")
.order_by(models.Metamessage.id.desc()) # get the most recent
.limit(1)
)
result = await db.execute(stmt)
representation = result.scalar_one_or_none()
return (
representation.content
if representation
else "No user representation available."
)
total_time = asyncio.get_event_loop().time() - stream_start
logger.debug(f"stream() setup completed in {total_time:.2f}s")
return stream
@observe()
@ -148,26 +151,115 @@ async def chat(
app_id: str,
user_id: str,
session_id: str,
query: schemas.AgentQuery,
queries: str | list[str],
stream: bool = False,
) -> schemas.AgentChat | MessageStreamManager:
questions = [query.queries] if isinstance(query.queries, str) else query.queries
) -> schemas.DialecticResponse | MessageStreamManager:
"""
Chat with the Dialectic API using on-demand user representation generation.
This function:
1. Sets up resources needed (embedding store, latest message ID)
2. Runs two parallel processes:
- Retrieves long-term facts from the vector store based on the query
- Gets recent chat history and runs ToM inference
3. Combines both into a fresh user representation
4. Uses this representation to answer the query
5. Saves the representation for future use
"""
# Format the query string
questions = [queries] if isinstance(queries, str) else queries
final_query = "\n".join(questions) if len(questions) > 1 else questions[0]
logger.debug(f"Received query: {final_query} for session {session_id}")
logger.debug("Starting on-demand user representation generation")
start_time = asyncio.get_event_loop().time()
async with SessionLocal() as db:
# Run user representation retrieval and chat history retrieval concurrently
user_rep_task = get_latest_user_representation(db, app_id, user_id)
history_task = chat_history(app_id, user_id, session_id)
# Setup phase - create resources we'll need for all operations
# 1. Create embedding store
collection = await crud.get_or_create_user_protected_collection(
db, app_id, user_id
)
embedding_store = CollectionEmbeddingStore(
db=db,
app_id=app_id,
user_id=user_id,
collection_id=collection.public_id, # type: ignore
)
logger.debug(
f"Created embedding store with collection_id: {collection.public_id if collection else None}"
)
# 2. Get the latest user message to attach the user representation to
stmt = (
select(models.Message)
.join(models.Session, models.Session.public_id == models.Message.session_id)
.join(models.User, models.User.public_id == models.Session.user_id)
.join(models.App, models.App.public_id == models.User.app_id)
.where(models.App.public_id == app_id)
.where(models.User.public_id == user_id)
.where(models.Message.session_id == session_id)
.where(models.Message.is_user)
.order_by(models.Message.id.desc())
.limit(1)
)
latest_messages = await db.execute(stmt)
latest_message = latest_messages.scalar_one_or_none()
latest_message_id = latest_message.public_id if latest_message else None
logger.debug(f"Latest user message ID: {latest_message_id}")
# Get chat history for the session
chat_history, _, _ = await history.get_summarized_history(
db, session_id, summary_type=history.SummaryType.SHORT
)
if not chat_history:
logger.warning(f"No chat history found for session {session_id}")
chat_history = f"someone asked this about the user's message: {final_query}"
logger.debug(f"IDs: {app_id}, {user_id}, {session_id}")
message_count = len(chat_history.split("\n"))
logger.debug(f"Retrieved chat history: {message_count} messages")
# Run both long-term and short-term context retrieval concurrently
logger.debug("Starting parallel tasks for context retrieval")
long_term_task = get_long_term_facts(final_query, embedding_store)
short_term_task = run_tom_inference(chat_history, session_id)
# Wait for both tasks to complete
user_representation, history = await asyncio.gather(user_rep_task, history_task)
facts, tom_inference = await asyncio.gather(long_term_task, short_term_task)
logger.debug(f"Retrieved {len(facts)} facts from long-term memory")
logger.debug(f"TOM inference completed with {len(tom_inference)} characters")
# Generate a fresh user representation
logger.debug("Generating user representation")
user_representation = await generate_user_representation(
app_id=app_id,
user_id=user_id,
session_id=session_id,
chat_history=chat_history,
tom_inference=tom_inference,
facts=facts,
embedding_store=embedding_store,
db=db,
message_id=latest_message_id,
with_inference=False,
)
logger.debug(
f"User representation generated: {len(user_representation)} characters"
)
# Create a Dialectic chain with the fresh user representation
chain = Dialectic(
agent_input=final_query,
user_representation=user_representation,
chat_history=history,
chat_history=chat_history,
)
generation_time = asyncio.get_event_loop().time() - start_time
logger.debug(f"User representation generation completed in {generation_time:.2f}s")
langfuse_context.update_current_trace(
session_id=session_id,
user_id=user_id,
@ -175,7 +267,285 @@ async def chat(
metadata={"environment": os.getenv("SENTRY_ENVIRONMENT")},
)
# Use streaming or non-streaming response based on the request
logger.debug(f"Calling Dialectic with streaming={stream}")
query_start_time = asyncio.get_event_loop().time()
if stream:
return chain.stream()
response = chain.call()
return schemas.AgentChat(content=response[0].text)
response_stream = await chain.stream()
logger.debug(
f"Dialectic stream started after {asyncio.get_event_loop().time() - query_start_time:.2f}s"
)
return response_stream
response = await chain.call()
query_time = asyncio.get_event_loop().time() - query_start_time
total_time = asyncio.get_event_loop().time() - start_time
logger.debug(
f"Dialectic response received in {query_time:.2f}s (total: {total_time:.2f}s)"
)
return schemas.DialecticResponse(content=response[0]["text"])
async def get_long_term_facts(
query: str, embedding_store: CollectionEmbeddingStore
) -> list[str]:
"""
Generate queries based on the dialectic query and retrieve relevant facts.
Args:
query: The user query
embedding_store: The embedding store to search
Returns:
List of retrieved facts
"""
logger.debug(f"Starting fact retrieval for query: {query}")
fact_start_time = asyncio.get_event_loop().time()
# Generate multiple queries for the semantic search
logger.debug("Generating semantic queries")
search_queries = await generate_semantic_queries(query)
logger.debug(f"Generated {len(search_queries)} semantic queries: {search_queries}")
# Create a list of coroutines, one for each query
async def execute_query(i: int, search_query: str) -> list[str]:
logger.debug(f"Starting query {i + 1}/{len(search_queries)}: {search_query}")
query_start = asyncio.get_event_loop().time()
facts = await embedding_store.get_relevant_facts(
search_query, top_k=10, max_distance=0.85
)
query_time = asyncio.get_event_loop().time() - query_start
logger.debug(f"Query {i + 1} retrieved {len(facts)} facts in {query_time:.2f}s")
return facts
# Execute all queries in parallel
query_tasks = [
execute_query(i, search_query) for i, search_query in enumerate(search_queries)
]
all_facts_lists = await asyncio.gather(*query_tasks)
# Combine all facts into a single set to remove duplicates
retrieved_facts = set()
for facts in all_facts_lists:
retrieved_facts.update(facts)
total_time = asyncio.get_event_loop().time() - fact_start_time
logger.debug(
f"Total fact retrieval completed in {total_time:.2f}s with {len(retrieved_facts)} unique facts"
)
return list(retrieved_facts)
async def run_tom_inference(chat_history: str, session_id: str) -> str:
"""
Run ToM inference on chat history.
Args:
chat_history: The chat history
session_id: The session ID
Returns:
The ToM inference
"""
# Run ToM inference
logger.debug(f"Running ToM inference for session {session_id}")
tom_start_time = asyncio.get_event_loop().time()
# Get chat history length to determine if this is a new conversation
tom_inference_response = await get_tom_inference(
chat_history, session_id, method="single_prompt", user_representation=""
)
# Extract the prediction from the response
tom_time = asyncio.get_event_loop().time() - tom_start_time
logger.debug(f"ToM inference completed in {tom_time:.2f}s")
prediction = parse_xml_content(tom_inference_response, "prediction")
logger.debug(f"Prediction length: {len(prediction)} characters")
return prediction
async def generate_semantic_queries(query: str) -> list[str]:
"""
Generate multiple semantically relevant queries based on the original query using LLM.
This helps retrieve more diverse and relevant facts from the vector store.
Args:
query: The original dialectic query
Returns:
A list of semantically relevant queries
"""
logger.debug(f"Generating semantic queries from: {query}")
query_start = asyncio.get_event_loop().time()
logger.debug("Calling LLM for query generation")
llm_start = asyncio.get_event_loop().time()
# Create a new model client
client = ModelClient(
provider=DEF_QUERY_GENERATION_PROVIDER, model=DEF_QUERY_GENERATION_MODEL
)
# Prepare the messages for Anthropic
messages: list[dict[str, Any]] = [{"role": "user", "content": query}]
# Generate the response
try:
result = await client.generate(
messages=messages,
system=QUERY_GENERATION_SYSTEM,
max_tokens=1000,
use_caching=True, # Likely not caching because the system prompt is under 1000 tokens
)
llm_time = asyncio.get_event_loop().time() - llm_start
logger.debug(f"LLM response received in {llm_time:.2f}s: {result[:100]}...")
# Parse the JSON response to get a list of queries
try:
queries = json.loads(result)
if not isinstance(queries, list):
# Fallback if response is not a valid list
logger.debug("LLM response not a list, using as single query")
queries = [result]
except json.JSONDecodeError:
# Fallback if response is not valid JSON
logger.debug("Failed to parse JSON response, using raw response as query")
queries = [query] # Fall back to the original query
# Ensure we always include the original query
if query not in queries:
logger.debug("Adding original query to results")
queries.append(query)
total_time = asyncio.get_event_loop().time() - query_start
logger.debug(f"Generated {len(queries)} queries in {total_time:.2f}s")
return queries
except Exception as e:
logger.error(f"Error during API call: {str(e)}")
raise
async def generate_user_representation(
app_id: str,
user_id: str,
session_id: str,
chat_history: str,
tom_inference: str,
facts: list[str],
embedding_store: CollectionEmbeddingStore,
db: AsyncSession,
message_id: Optional[str] = None,
with_inference: bool = False,
) -> str:
"""
Generate a user representation by combining long-term facts and short-term context.
Optionally save it as a metamessage if message_id is provided.
Only uses existing representations from the same session for continuity.
Returns:
The generated user representation.
"""
logger.debug("Starting user representation generation")
rep_start_time = asyncio.get_event_loop().time()
if with_inference:
# Fetch the latest user representation from the same session
logger.debug(f"Fetching latest representation for session {session_id}")
latest_representation_stmt = (
select(models.Metamessage)
.join(
models.Message,
models.Message.public_id == models.Metamessage.message_id,
)
.join(models.Session, models.Message.session_id == models.Session.public_id)
.where(models.Session.public_id == session_id) # Only from the same session
.where(
models.Metamessage.metamessage_type
== USER_REPRESENTATION_METAMESSAGE_TYPE
)
.order_by(models.Metamessage.id.desc())
.limit(1)
)
result = await db.execute(latest_representation_stmt)
latest_representation_obj = result.scalar_one_or_none()
latest_representation = (
latest_representation_obj.content
if latest_representation_obj
else "No user representation available."
)
logger.debug(
f"Found previous representation: {len(latest_representation)} characters"
)
logger.debug(f"Using {len(facts)} facts for representation")
# Generate the new user representation
logger.debug("Calling get_user_representation")
gen_start_time = asyncio.get_event_loop().time()
user_representation_response = await get_user_representation_long_term(
chat_history=chat_history,
session_id=session_id,
facts=facts,
embedding_store=embedding_store,
user_representation=latest_representation,
tom_inference=tom_inference,
)
gen_time = asyncio.get_event_loop().time() - gen_start_time
logger.debug(f"get_user_representation completed in {gen_time:.2f}s")
# Extract the representation from the response
representation = parse_xml_content(
user_representation_response, "representation"
)
logger.debug(f"Extracted representation: {len(representation)} characters")
else:
representation = f"""
PREDICTION ABOUT THE USER'S CURRENT MENTAL STATE:
{tom_inference}
RELEVANT LONG-TERM FACTS ABOUT THE USER:
{facts}
"""
logger.debug(f"Representation: {representation}")
# If message_id is provided, save the representation as a metamessage
if not representation:
logger.debug("Empty representation, skipping save")
else:
logger.debug(f"Saving representation to message_id: {message_id}")
save_start = asyncio.get_event_loop().time()
try:
async with SessionLocal() as save_db:
try:
# First check if message exists
message_check_stmt = select(models.Message).where(
models.Message.public_id == message_id
)
message_check = await save_db.execute(message_check_stmt)
message_exists = message_check.scalar_one_or_none() is not None
if not message_exists:
message_id = None
else:
metamessage = models.Metamessage(
user_id=user_id,
session_id=session_id,
message_id=message_id if message_id else None,
metamessage_type=USER_REPRESENTATION_METAMESSAGE_TYPE,
content=representation,
h_metadata={},
)
save_db.add(metamessage)
await save_db.commit()
save_time = asyncio.get_event_loop().time() - save_start
logger.debug(f"Representation saved in {save_time:.2f}s")
except Exception as inner_e:
logger.error(f"Error during save DB operation: {str(inner_e)}")
await save_db.rollback()
except Exception as e:
logger.error(f"Error creating DB session: {str(e)}")
total_time = asyncio.get_event_loop().time() - rep_start_time
logger.debug(f"Total representation generation completed in {total_time:.2f}s")
return representation

651
src/crud.py
View File

File diff suppressed because it is too large Load Diff

25
src/db.py
View File

@ -1,9 +1,7 @@
import os
from alembic import command
from alembic.config import Config
from dotenv import load_dotenv
from sqlalchemy import MetaData, create_engine, inspect
from sqlalchemy import MetaData, create_engine, inspect, text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.orm import declarative_base
@ -13,11 +11,6 @@ connect_args = {
"prepare_threshold": None,
}
# if (
# os.environ["DATABASE_TYPE"] == "sqlite"
# ): # https://fastapi.tiangolo.com/tutorial/sql-databases/#note
# connect_args = {"check_same_thread": False}
engine = create_async_engine(
os.environ["CONNECTION_URI"],
connect_args=connect_args,
@ -48,6 +41,12 @@ def scaffold_db():
"""use a sync engine for scaffolding the database. ddl operations are unavailable
with async engines
"""
# Debug: Print all tables that should be created
print("Tables defined in Base.metadata:")
for table in Base.metadata.sorted_tables:
print(f" - {table.name}")
# Create engine
engine = create_engine(
os.environ["CONNECTION_URI"],
@ -58,6 +57,11 @@ def scaffold_db():
# Create inspector to check if database exists
inspector = inspect(engine)
if table_schema:
with engine.connect() as connection:
connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{table_schema}"'))
connection.commit()
print(inspector.get_table_names(Base.metadata.schema))
# If no tables exist, create them with SQLAlchemy
@ -67,8 +71,3 @@ def scaffold_db():
# Clean up
engine.dispose()
# Run Alembic migrations regardless
print("Running database migrations...")
alembic_cfg = Config("alembic.ini")
command.upgrade(alembic_cfg, "head")

8
src/deriver/__main__.py
View File

@ -5,8 +5,14 @@ import uvloop
from .queue import main
if __name__ == "__main__":
print("[DERIVER] Starting deriver queue processor")
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
try:
print("[DERIVER] Running main loop")
asyncio.run(main())
except KeyboardInterrupt:
print("Shutdown initiated via KeyboardInterrupt")
print("[DERIVER] Shutdown initiated via KeyboardInterrupt")
except Exception as e:
print(f"[DERIVER] Error in main process: {str(e)}")
finally:
print("[DERIVER] Deriver process exiting")

341
src/deriver/consumer.py
View File

@ -1,113 +1,63 @@
import logging
import os
import re
import sentry_sdk
from langfuse.decorators import langfuse_context, observe
from langfuse.decorators import observe
from rich.console import Console
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from .. import models
from ..exceptions import ResourceNotFoundException, ValidationException
from .tom import get_tom_inference, get_user_representation
from .. import crud
from ..utils import history
from .tom.embeddings import CollectionEmbeddingStore
from .tom.long_term import extract_facts_long_term
# Configure logging
logger = logging.getLogger(__name__)
# Turn off SQLAlchemy Echo logging
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
console = Console(markup=False)
TOM_METHOD = os.getenv("TOM_METHOD", "single_prompt")
USER_REPRESENTATION_METHOD = os.getenv("USER_REPRESENTATION_METHOD", "single_prompt")
USER_REPRESENTATION_METHOD = os.getenv("USER_REPRESENTATION_METHOD", "long_term")
# FIXME see if this is SAFE
async def add_metamessage(db, message_id, metamessage_type, content):
metamessage = models.Metamessage(
message_id=message_id,
metamessage_type=metamessage_type,
content=content,
h_metadata={},
)
db.add(metamessage)
def parse_xml_content(text, tag):
pattern = f"<{tag}>(.*?)</{tag}>"
match = re.search(pattern, text, re.DOTALL)
return match.group(1).strip() if match else ""
async def get_chat_history(db, session_id, message_id) -> str:
subquery = (
select(models.Message.id)
.where(models.Message.public_id == message_id)
.scalar_subquery()
)
messages_stmt = (
select(models.Message)
.where(models.Message.session_id == session_id)
.order_by(models.Message.id.desc())
.where(models.Message.id < subquery)
.limit(10)
)
result = await db.execute(messages_stmt)
messages = result.scalars().all()[::-1]
chat_history_str = "\n".join(
[f"human: {m.content}" if m.is_user else f"ai: {m.content}" for m in messages]
)
return chat_history_str
# async def add_metamessage(db, message_id, metamessage_type, content):
# metamessage = models.Metamessage(
# message_id=message_id,
# metamessage_type=metamessage_type,
# content=content,
# h_metadata={},
# )
# db.add(metamessage)
async def process_item(db: AsyncSession, payload: dict):
"""
Process a queue item based on whether it's a user or AI message.
Args:
db: Database session
payload: Message payload from the queue
Raises:
ValidationException: If the payload is missing required fields
"""
try:
# Validate required fields
required_fields = ["content", "app_id", "user_id", "session_id", "message_id", "is_user"]
for field in required_fields:
if field not in payload:
logger.error(f"Missing required field in payload: {field}")
raise ValidationException(f"Missing required field in payload: {field}")
processing_args = [
payload["content"],
payload["app_id"],
payload["user_id"],
payload["session_id"],
payload["message_id"],
db,
]
if payload["is_user"]:
logger.info(f"Processing user message: {payload['message_id']}")
await process_user_message(*processing_args)
else:
logger.info(f"Processing AI message: {payload['message_id']}")
await process_ai_message(*processing_args)
except Exception as e:
logger.error(f"Error processing message {payload.get('message_id', 'unknown')}: {str(e)}")
if os.getenv("SENTRY_ENABLED", "False").lower() == "true":
sentry_sdk.capture_exception(e)
raise
logger.debug(
f"process_item received payload: {payload['message_id']} is_user={payload['is_user']}"
)
processing_args = [
payload["content"],
payload["app_id"],
payload["user_id"],
payload["session_id"],
payload["message_id"],
db,
]
if payload["is_user"]:
logger.debug(f"Processing user message: {payload['message_id']}")
await process_user_message(*processing_args)
else:
logger.debug(f"Processing AI message: {payload['message_id']}")
await process_ai_message(*processing_args)
logger.debug(f"Finished processing message: {payload['message_id']}")
await summarize_if_needed(
db, payload["session_id"], payload["user_id"], payload["message_id"]
)
return
@sentry_sdk.trace
@observe()
# @observe()
async def process_ai_message(
content: str,
app_id: str,
@ -133,89 +83,162 @@ async def process_user_message(
db: AsyncSession,
):
"""
Process a user message by:
- Getting TOM inference
- Getting user representation
Process a user message by extracting facts and saving them to the vector store.
This runs as a background process after a user message is logged.
"""
console.print(f"Processing User Message: {content}", style="orange1")
process_start = os.times()[4] # Get current CPU time
logger.debug(f"Starting fact extraction for user message: {message_id}")
# Get chat history and append current message
chat_history_str = await get_chat_history(db, session_id, message_id)
chat_history_str = f"{chat_history_str}\nhuman: {content}"
# Get TOM inference, parse and save it
tom_inference_response = await get_tom_inference(
chat_history_str, session_id, method=TOM_METHOD
logger.debug(f"Retrieving chat history for session: {session_id}")
(
short_history_text,
short_history_messages,
latest_short_summary,
) = await history.get_summarized_history(
db, session_id, summary_type=history.SummaryType.SHORT
)
tom_inference = parse_xml_content(tom_inference_response, "prediction")
await add_metamessage(
db,
message_id,
"tom_inference",
tom_inference,
chat_history_str = f"{short_history_text}\nhuman: {content}"
# Extract facts from chat history
logger.debug("Extracting facts from chat history")
extract_start = os.times()[4]
facts = await extract_facts_long_term(chat_history_str)
extract_time = os.times()[4] - extract_start
console.print(f"Extracted Facts: {facts}", style="bright_blue")
logger.debug(f"Extracted {len(facts)} facts in {extract_time:.2f}s")
# Save the facts to the collection
logger.debug(f"Setting up embedding store for app: {app_id}, user: {user_id}")
collection = await crud.get_or_create_user_protected_collection(
db=db, app_id=app_id, user_id=user_id
)
await db.commit()
# Fetch the latest user representation
user_representation_stmt = (
select(models.Metamessage)
.join(
models.Message,
models.Message.public_id == models.Metamessage.message_id,
)
.join(
models.Session,
models.Message.session_id == models.Session.public_id,
)
.join(models.User, models.User.public_id == models.Session.user_id)
.join(models.App, models.App.public_id == models.User.app_id)
.where(models.App.public_id == app_id)
.where(models.User.public_id == user_id)
.where(models.Metamessage.metamessage_type == "user_representation")
.order_by(models.Metamessage.id.desc()) # get the most recent
.limit(1)
)
response = await db.execute(user_representation_stmt)
existing_representation = response.scalar_one_or_none()
existing_representation_content = (
existing_representation.content if existing_representation else "None"
)
logger.info(f"User {user_id}: Existing Representation retrieved")
logger.debug(f"User {user_id}: Existing Representation: {existing_representation_content}")
langfuse_context.update_current_trace(
session_id=session_id,
embedding_store = CollectionEmbeddingStore(
db=db,
app_id=app_id,
user_id=user_id,
release=os.getenv("SENTRY_RELEASE"),
metadata={"environment": os.getenv("SENTRY_ENVIRONMENT")},
collection_id=collection.public_id, # type: ignore
)
# Call user_representation
user_representation_response = await get_user_representation(
chat_history=chat_history_str,
session_id=session_id,
user_representation=existing_representation_content,
tom_inference=tom_inference,
method=USER_REPRESENTATION_METHOD,
# Filter out facts that are duplicates of existing facts in the vector store
logger.debug("Removing duplicate facts")
dedup_start = os.times()[4]
unique_facts = await embedding_store.remove_duplicates(facts)
dedup_time = os.times()[4] - dedup_start
logger.debug(
f"Found {len(unique_facts)}/{len(facts)} unique facts in {dedup_time:.2f}s"
)
# parse the user_representation response
user_representation_response = parse_xml_content(
user_representation_response, "representation"
# Only save the unique facts
if unique_facts:
logger.debug(f"Saving {len(unique_facts)} unique facts to vector store")
save_start = os.times()[4]
await embedding_store.save_facts(unique_facts, message_id=message_id)
save_time = os.times()[4] - save_start
logger.debug(f"Facts saved in {save_time:.2f}s")
else:
logger.debug("No unique facts to save")
console.print(f"Saved {len(unique_facts)} unique facts", style="bright_green")
total_time = os.times()[4] - process_start
logger.debug(f"Total processing time: {total_time:.2f}s")
async def summarize_if_needed(
db: AsyncSession, session_id: str, user_id: str, message_id: str
):
summary_start = os.times()[4]
logger.debug("Checking if summaries should be created")
# STEP 1: First check if we need a short summary (every 10 messages)
(
should_create_short,
short_messages,
latest_short_summary,
) = await history.should_create_summary(
db, session_id, summary_type=history.SummaryType.SHORT
)
# Store the user_representation response as a metamessage
await add_metamessage(
db,
message_id,
"user_representation",
user_representation_response,
)
await db.commit()
if should_create_short:
logger.debug(f"Short summary needed for {len(short_messages)} messages")
console.print(
f"User Representation:\n{user_representation_response}",
style="bright_green",
)
# STEP 2: If we need a short summary, check if we also need a long summary
(
should_create_long,
long_messages,
latest_long_summary,
) = await history.should_create_summary(
db, session_id, summary_type=history.SummaryType.LONG
)
# STEP 3: If we need a long summary, create it first before creating the short summary
if should_create_long:
logger.debug(
f"Creating new long summary covering {len(long_messages)} messages"
)
try:
# Get previous long summary context if available
previous_long_summary = (
latest_long_summary.content if latest_long_summary else None
)
# Create a new long summary
long_summary_text = await history.create_summary(
messages=long_messages,
previous_summary=previous_long_summary,
summary_type=history.SummaryType.LONG,
)
# Save the long summary as a metamessage and capture the returned object
latest_long_summary = await history.save_summary_metamessage(
db=db,
user_id=user_id,
session_id=session_id,
message_id=message_id,
summary_content=long_summary_text,
message_count=len(long_messages),
summary_type=history.SummaryType.LONG,
)
logger.debug("Long summary created and saved successfully")
except Exception as e:
logger.error(f"Error creating long summary: {str(e)}")
else:
logger.debug(
f"No long summary needed. Need {history.MESSAGES_PER_LONG_SUMMARY} messages since last long summary."
)
# STEP 4: Now create the short summary, using the latest long summary for context if available
logger.debug(
f"Creating new short summary covering {len(short_messages)} messages"
)
try:
previous_summary = (
latest_long_summary.content if latest_long_summary else None
)
# Create a new short summary
short_summary_text = await history.create_summary(
messages=short_messages,
previous_summary=previous_summary,
summary_type=history.SummaryType.SHORT,
)
# Save the short summary as a metamessage
await history.save_summary_metamessage(
db=db,
user_id=user_id,
session_id=session_id,
message_id=message_id,
summary_content=short_summary_text,
message_count=len(short_messages),
summary_type=history.SummaryType.SHORT,
)
logger.debug("Short summary created and saved successfully")
except Exception as e:
logger.error(f"Error creating short summary: {str(e)}")
else:
logger.debug(
f"No short summary needed. Need {history.MESSAGES_PER_SHORT_SUMMARY} messages since last short summary."
)
summary_time = os.times()[4] - summary_start
logger.debug(f"Summary check completed in {summary_time:.2f}s")

71
src/deriver/queue.py
View File

@ -1,8 +1,8 @@
import asyncio
import logging
import os
import signal
from datetime import datetime, timedelta
from logging import getLogger
import sentry_sdk
from dotenv import load_dotenv
@ -16,7 +16,7 @@ from .. import models
from ..db import SessionLocal
from .consumer import process_item
logger = logging.getLogger(__name__)
logger = getLogger(__name__)
load_dotenv()
@ -57,13 +57,19 @@ class QueueManager:
async def initialize(self):
"""Setup signal handlers and start the main polling loop"""
logger.debug(f"Initializing QueueManager with {self.workers} workers")
# Set up signal handlers
loop = asyncio.get_running_loop()
signals = (signal.SIGTERM, signal.SIGINT)
for sig in signals:
loop.add_signal_handler(
sig, lambda s=sig: asyncio.create_task(self.shutdown(s))
)
logger.debug("Signal handlers registered")
# Run the polling loop directly in this task
logger.debug("Starting polling loop directly")
try:
await self.polling_loop()
finally:
@ -75,7 +81,9 @@ class QueueManager:
self.shutdown_event.set()
if self.active_tasks:
logger.info(f"Waiting for {len(self.active_tasks)} active tasks to complete...")
logger.info(
f"Waiting for {len(self.active_tasks)} active tasks to complete..."
)
await asyncio.gather(*self.active_tasks, return_exceptions=True)
async def cleanup(self):
@ -86,7 +94,9 @@ class QueueManager:
async with SessionLocal() as db:
await db.execute(
delete(models.ActiveQueueSession).where(
models.ActiveQueueSession.session_id.in_(self.owned_sessions)
models.ActiveQueueSession.session_id.in_(
self.owned_sessions
)
)
)
await db.commit()
@ -128,15 +138,18 @@ class QueueManager:
async def polling_loop(self):
"""Main polling loop to find and process new sessions"""
logger.debug("Starting polling loop")
try:
while not self.shutdown_event.is_set():
if self.queue_empty_flag.is_set():
# logger.debug("Queue empty flag set, waiting")
await asyncio.sleep(1)
self.queue_empty_flag.clear()
continue
# Chec if we have capacity before querying
# Check if we have capacity before querying
if self.semaphore.locked():
# logger.debug("All workers busy, waiting")
await asyncio.sleep(1) # Wait before trying again
continue
@ -150,13 +163,16 @@ class QueueManager:
# Try to claim the session
await db.execute(
insert(models.ActiveQueueSession).values(
session_id=session_id
session_id=session_id,
)
)
await db.commit()
# Track this session
self.track_session(session_id)
logger.debug(
f"Claimed session {session_id} for processing"
)
# Create a new task for processing this session
if not self.shutdown_event.is_set():
@ -166,6 +182,9 @@ class QueueManager:
self.add_task(task)
except IntegrityError:
await db.rollback()
logger.debug(
f"Failed to claim session {session_id}, already owned"
)
else:
self.queue_empty_flag.set()
await asyncio.sleep(1)
@ -185,39 +204,60 @@ class QueueManager:
@sentry_sdk.trace
async def process_session(self, session_id: int):
"""Process all messages for a session"""
logger.debug(f"Starting to process session {session_id}")
async with self.semaphore: # Hold the semaphore for the entire session duration
async with SessionLocal() as db:
try:
message_count = 0
while not self.shutdown_event.is_set():
message = await self.get_next_message(db, session_id)
if not message:
logger.debug(f"No more messages for session {session_id}")
break
message_count += 1
logger.debug(
f"Processing message {message.id} for session {session_id} (message {message_count})"
)
try:
logger.info(f"Processing message {message.id} from session {session_id}")
logger.info(
f"Processing message {message.id} from session {session_id}"
)
await process_item(db, payload=message.payload)
logger.info(f"Successfully processed message {message.id}")
logger.debug(f"Successfully processed message {message.id}")
except Exception as e:
logger.error(f"Error processing message {message.id}: {str(e)}", exc_info=True)
logger.error(
f"Error processing message {message.id}: {str(e)}",
exc_info=True,
)
if os.getenv("SENTRY_ENABLED", "False").lower() == "true":
sentry_sdk.capture_exception(e)
finally:
# Prevent malformed messages from stalling queue indefinitely
message.processed = True
await db.commit()
logger.info(f"Marked message {message.id} as processed")
logger.debug(f"Marked message {message.id} as processed")
if self.shutdown_event.is_set():
logger.debug(
f"Shutdown requested, stopping processing for session {session_id}"
)
break
# Update last_updated timestamp to showthis session is still being processed
# Update last_updated timestamp to show this session is still being processed
await db.execute(
update(models.ActiveQueueSession)
.where(models.ActiveQueueSession.session_id == session_id)
.values(last_updated=func.now())
)
await db.commit()
logger.debug(
f"Completed processing session {session_id}, processed {message_count} messages"
)
finally:
# Remove session from active_sessions when done
logger.debug(f"Removing session {session_id} from active sessions")
await db.execute(
delete(models.ActiveQueueSession).where(
models.ActiveQueueSession.session_id == session_id
@ -241,5 +281,12 @@ class QueueManager:
async def main():
logger.debug("Starting queue manager")
manager = QueueManager()
await manager.initialize()
try:
await manager.initialize()
except Exception as e:
logger.error(f"Error in main: {str(e)}")
sentry_sdk.capture_exception(e)
finally:
logger.debug("Main function exiting")

13
src/deriver/tom/README.md Normal file
View File

@ -0,0 +1,13 @@
# Theory of Mind Inference
[Theory of Mind](https://blog.plasticlabs.ai/blog/Theory-of-Mind-Is-All-You-Need) is a core principle behind Honcho: we believe that enabling AI agents to reason about users' mental states is essential if we want them to successfully act on our behalf.
Honcho currently features three different modules for theory of mind inference:
- `conversational.py`: Inspired by our work on [metanarrative prompting](https://blog.plasticlabs.ai/blog/Agent-Identity). Uses a metanarrative prompt for both ToM inference and generating a user representation.
- `single_prompt.py`: A more conventional and straightforward approach that specifies in a single system prompt what it wants the LLM to output.
- `long_term.py`: Formats a theory of mind inference and a series of long-term facts into a user representation.
The current setup works as follows:
- We extract facts from incoming messages using the code in `src.deriver.consumer`.
- These messages get added to the protected `honcho` user collection using the `CollectionEmbeddingStore` in `src.deriver.tom.embeddings`.
- The dialectic endpoint, in `src.agent`, retrieves long-term facts from this store that are relevant to the query, and runs the ToM inference in `src.deriver.tom.single_prompt` to generate a prediction of the user's short-term mental state.
- The retrieved long-term facts and the short-term ToM inference are combined into a user representation. By default, this is done using a simple f-string, but they can optionally be combined using a separate inference, which would use `src.deriver.tom.long_term`.

62
src/deriver/tom/__init__.py
View File

@ -1,30 +1,52 @@
from .conversational import get_tom_inference_conversational, get_user_representation_conversational
from .single_prompt import get_tom_inference_single_prompt, get_user_representation_single_prompt
from .conversational import (
get_tom_inference_conversational,
get_user_representation_conversational,
)
from .long_term import get_user_representation_long_term
from .single_prompt import (
get_tom_inference_single_prompt,
get_user_representation_single_prompt,
)
async def get_tom_inference(chat_history: str,
session_id: str,
user_representation: str = "None",
method: str = "conversational",
**kwargs
) -> str:
async def get_tom_inference(
chat_history: str,
session_id: str,
user_representation: str = "None",
method: str = "conversational",
**kwargs,
) -> str:
if method == "conversational":
return await get_tom_inference_conversational(chat_history, session_id, user_representation, **kwargs)
return await get_tom_inference_conversational(
chat_history, session_id, user_representation, **kwargs
)
elif method == "single_prompt":
return await get_tom_inference_single_prompt(chat_history, session_id, user_representation, **kwargs)
return await get_tom_inference_single_prompt(
chat_history, session_id, user_representation, **kwargs
)
else:
raise ValueError(f"Invalid method: {method}")
async def get_user_representation(chat_history: str,
session_id: str,
user_representation: str = "None",
tom_inference: str = "None",
method: str = "conversational",
**kwargs
) -> str:
async def get_user_representation(
chat_history: str,
session_id: str,
user_representation: str = "None",
tom_inference: str = "None",
method: str = "conversational",
**kwargs,
) -> str:
if method == "conversational":
return await get_user_representation_conversational(chat_history, session_id, user_representation, tom_inference, **kwargs)
return await get_user_representation_conversational(
chat_history, session_id, user_representation, tom_inference, **kwargs
)
elif method == "single_prompt":
return await get_user_representation_single_prompt(chat_history, session_id, user_representation, tom_inference, **kwargs)
return await get_user_representation_single_prompt(
chat_history, session_id, user_representation, tom_inference, **kwargs
)
elif method == "long_term":
return await get_user_representation_long_term(
chat_history, session_id, user_representation, tom_inference, **kwargs
)
else:
raise ValueError(f"Invalid method: {method}")
raise ValueError(f"Invalid method: {method}")

4
src/deriver/tom/conversational.py
View File

@ -15,7 +15,7 @@ anthropic = Anthropic(
@ai_track("Tom Inference")
@observe(as_type="generation")
@observe()
async def get_tom_inference_conversational(
chat_history: str, session_id: str, user_representation: str = "None"
) -> str:
@ -81,7 +81,7 @@ async def get_tom_inference_conversational(
@ai_track("User Representation")
@observe(as_type="generation")
@observe()
async def get_user_representation_conversational(
chat_history: str,
session_id: str,

114
src/deriver/tom/embeddings.py Normal file
View File

@ -0,0 +1,114 @@
import logging
from sqlalchemy.ext.asyncio import AsyncSession
from ... import crud, schemas
logger = logging.getLogger(__name__)
class CollectionEmbeddingStore:
def __init__(self, db: AsyncSession, app_id: str, user_id: str, collection_id: str):
self.db = db
self.app_id = app_id
self.user_id = user_id
self.collection_id = collection_id
async def save_facts(
self,
facts: list[str],
replace_duplicates: bool = True,
similarity_threshold: float = 0.85,
message_id: str = None,
) -> None:
"""Save facts to the collection.
Args:
facts: List of facts to save
replace_duplicates: If True, replace old duplicates with new facts. If False, discard new duplicates
similarity_threshold: Facts with similarity above this threshold are considered duplicates
"""
for fact in facts:
# Create document with duplicate checking
try:
metadata = {}
if message_id:
metadata["message_id"] = message_id
await crud.create_document(
self.db,
document=schemas.DocumentCreate(content=fact, metadata=metadata),
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
duplicate_threshold=1
- similarity_threshold, # Convert similarity to distance
)
except Exception as e:
logger.error(f"Error creating document: {e}")
continue
async def get_relevant_facts(
self, query: str, top_k: int = 5, max_distance: float = 0.3
) -> list[str]:
"""Retrieve the most relevant facts for a given query.
Args:
query: The query text to find relevant facts for
top_k: Maximum number of facts to return
similarity_threshold: Minimum similarity score for a fact to be considered relevant
Returns:
List of facts sorted by relevance
"""
documents = await crud.query_documents(
self.db,
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
query=query,
max_distance=max_distance,
top_k=top_k,
)
return [doc.content for doc in documents]
async def remove_duplicates(
self, facts: list[str], similarity_threshold: float = 0.85
) -> list[str]:
"""Remove facts that are duplicates of existing facts in the vector store.
Args:
facts: List of facts to check for duplicates
similarity_threshold: Facts with similarity above this threshold are considered duplicates
Returns:
List of facts that are not duplicates of existing facts
"""
unique_facts = []
for fact in facts:
try:
# Check for duplicates using the crud function
duplicates = await crud.get_duplicate_documents(
self.db,
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
content=fact,
similarity_threshold=similarity_threshold,
)
if not duplicates:
# No duplicates found, add to unique facts
unique_facts.append(fact)
else:
# Log duplicate found
logger.debug(
f"Duplicate found: {duplicates[0].content}. Ignoring fact: {fact}"
)
except Exception as e:
logger.error(f"Error checking for duplicates: {e}")
# If there's an error, still include the fact to avoid losing information
unique_facts.append(fact)
return unique_facts

216
src/deriver/tom/long_term.py Normal file
View File

@ -0,0 +1,216 @@
import json
import logging
import time
from typing import Optional
from langfuse.decorators import observe
from sentry_sdk.ai.monitoring import ai_track
from src.utils import parse_xml_content
from src.utils.model_client import ModelClient, ModelProvider
from .embeddings import CollectionEmbeddingStore
# Configure logging
logger = logging.getLogger(__name__)
# Constants for fact extraction
FACT_EXTRACTION_PROVIDER = ModelProvider.GEMINI
FACT_EXTRACTION_MODEL = "gemini-2.0-flash-lite"
USER_REPRESENTATION_PROVIDER = ModelProvider.GROQ
USER_REPRESENTATION_MODEL = "llama-3.3-70b-versatile"
MAX_FACT_DISTANCE = 0.85
@ai_track("User Representation")
@observe()
async def get_user_representation_long_term(
chat_history: str,
session_id: str,
embedding_store: CollectionEmbeddingStore,
user_representation: str = "None",
tom_inference: str = "None",
facts: Optional[list[str]] = None,
) -> str:
if facts is None:
facts = []
facts_str = "\n".join([f"- {fact}" for fact in facts])
logger.debug(f"Facts: {facts_str}")
system_prompt = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
Your job is to update the existing user representation (if provided) with the new information from the conversation history and theory of mind analysis.
REQUIREMENTS:
1. Distinguish between temporary states and persistent patterns
2. Only incorporate verified information into core profile
3. Track certainty levels for all information
4. Maintain areas of uncertainty explicitly
5. Update representation incrementally
6. DO NOT generate persistent information - it will be injected separately. Always include the <KNOWN_FACTS> tag in your response in order to inject the facts.
OUTPUT FORMAT:
<representation>
CURRENT STATE:
- Active Context: Current situation/activity
- Temporary Conditions: Immediate circumstances
<CURRENT_CURSOR_POSITION>
- Present Mood/Activity: What user is doing right now
<KNOWN_FACTS>
TENTATIVE PATTERNS:
- Possible Traits: Mark confidence (Low/Medium/High)
- Potential Interests: Need more evidence
- Speculative Elements: Clearly marked as unconfirmed
KNOWLEDGE GAPS:
- List key missing information
- Note areas needing clarification
EXPECTATION VIOLATIONS:
- Based on the above information, if the next message were to surprise you, what could it contain?
- Format: "POTENTIAL SURPRISE: [possible content] [reason] [confidence level]"
- Include 3-5 possible surprises
UPDATES:
- New Information: Recent observations
- Changes: Modified interpretations
- Removals: Information no longer supported
</representation>
"""
# Build the context message
context_str = f"CONVERSATION:\n{chat_history}\n\n"
if tom_inference != "None":
context_str += f"PREDICTION OF USER MENTAL STATE - MIGHT BE INCORRECT:\n{tom_inference}\n\n"
if user_representation != "None":
context_str += f"EXISTING USER REPRESENTATION - INCOMPLETE, TO BE UPDATED:\n{user_representation}"
messages = [
{
"role": "user",
"content": f"Please analyze this information and provide an updated user representation. DO NOT generate persistent information - it will be injected separately:\n{context_str}",
}
]
# Create a new model client
client = ModelClient(
provider=USER_REPRESENTATION_PROVIDER, model=USER_REPRESENTATION_MODEL
)
# Generate the response with caching enabled
response = await client.generate(
messages=messages,
system=system_prompt,
max_tokens=1000,
temperature=0,
use_caching=True, # Enable caching for the system prompt
)
# Inject the facts into the response
persistent_info = f"""PERSISTENT INFORMATION:
{facts_str}"""
return response.replace("<KNOWN_FACTS>", persistent_info)
@ai_track("Fact Extraction")
@observe()
async def extract_facts_long_term(chat_history: str) -> list[str]:
logger.debug("Starting fact extraction from chat history")
extract_start = time.time()
system_prompt = """
You are an AI assistant specialized in extracting and formatting relevant information about users from conversations. Your task is to analyze a given conversation and create a list of concise, factual statements about the user. These statements will be stored in a vector embedding database to enhance future interactions.
Here is the conversation you need to analyze:
<conversation>
{chat_history}
</conversation>
Instructions:
1. Carefully read through the conversation. Extract only new facts, from only the last message sent by the user - treat the rest of the conversation only as context. Ignore facts in the last message that are already stated in the conversation.
2. Identify key new pieces of information from the last message sent by the user that would be valuable for future interactions. Look for:
- Personal details (name, age, occupation, location, etc.)
- Preferences (likes, dislikes, interests, hobbies)
- Experiences (travel, education, work history)
- Expressive style (writing style, tone, etc.)
- Relationships (family, friends, pets)
- Goals or aspirations
- Challenges or problems they're facing
- Opinions or beliefs
3. For each piece of information you identify:
a. Verify that it is factual and explicitly stated in the conversation, not inferred.
b. Formulate it as a concise statement that would aid in semantic retrieval.
c. Ensure it is not similar to information previously stated in the conversation.
4. Before providing your final output, wrap your analysis in <information_extraction> tags. In this analysis:
- List each piece of information you've identified.
- For each piece of information:
* Quote the relevant part of the conversation.
* Categorize the information (e.g., personal detail, preference, experience).
* Explain why you've included this information.
* Show how you've formulated the fact for optimal semantic retrieval.
- Discuss any challenges you encountered in extracting or formatting the information.
5. After your analysis, provide your final output as a JSON array of strings. Each string should be a single fact about the user. Wrap the facts in <facts> tags.
Example of the expected output format:
<information_extraction>
[Analysis goes here]
</information_extraction>
<facts>
{{
"facts":
[
"User is 28 years old",
"User's friend Mary works as a software engineer",
"Favorite food is sushi"
]
}}
</facts>
Remember to focus on clear, concise statements that capture key information about the user. Each fact should be worded in a way that will aid its semantic retrieval from a vector embedding database. It's OK for this section to be quite long.
"""
message = system_prompt.format(chat_history=chat_history)
messages = [{"role": "user", "content": message}]
logger.debug("Calling LLM for fact extraction")
llm_start = time.time()
# Create a new model client
client = ModelClient(provider=FACT_EXTRACTION_PROVIDER, model=FACT_EXTRACTION_MODEL)
# Generate the response with caching enabled
response = await client.generate(
messages=messages,
max_tokens=1000,
temperature=0.0,
use_caching=True, # Enable caching for the system prompt
)
llm_time = time.time() - llm_start
logger.debug(f"LLM response received in {llm_time:.2f}s")
try:
logger.debug("Parsing JSON response")
facts_str = parse_xml_content(response, "facts")
response_data = json.loads(facts_str)
facts = response_data["facts"]
logger.debug(f"Extracted {len(facts)} facts")
if facts:
logger.debug(f"Sample facts: {facts[:3] if len(facts) > 3 else facts}")
except (json.JSONDecodeError, KeyError) as e:
logger.error(f"Error parsing response: {str(e)}")
facts = []
total_time = time.time() - extract_start
logger.debug(f"Total extraction completed in {total_time:.2f}s")
return facts

179
src/deriver/tom/single_prompt.py
View File

@ -1,28 +1,18 @@
import os
import logging
from typing import Any, Optional
import sentry_sdk
from anthropic import Anthropic
from langfuse.decorators import langfuse_context, observe
from sentry_sdk.ai.monitoring import ai_track
# Place the code below at the beginning of your application to initialize the tracer
from src.utils.model_client import ModelClient, ModelProvider
# Initialize the Anthropic client
anthropic = Anthropic(
api_key=os.getenv("ANTHROPIC_API_KEY"),
max_retries=5,
)
logger = logging.getLogger(__name__)
ANTHROPIC_MODEL = "claude-3-5-haiku-20241022"
DEF_PROVIDER = ModelProvider.GROQ
DEF_MODEL = "llama-3.3-70b-versatile"
@ai_track("Tom Inference")
@observe(as_type="generation")
async def get_tom_inference_single_prompt(
chat_history: str, session_id: str, user_representation: str = "None", **kwargs
) -> str:
with sentry_sdk.start_transaction(op="tom-inference", name="ToM Inference"):
system_prompt = """You are a system for analyzing conversations to make evidence-based inferences about user mental states.
TOM_SYSTEM_PROMPT = """You are a system for analyzing conversations to make evidence-based inferences about user mental states.
REQUIREMENTS:
1. Only make inferences that are directly supported by conversation evidence
@ -56,53 +46,9 @@ EXPECTATION VIOLATIONS:
- Based on the above information, if the next message were to surprise you, what could it contain?
- Format: "POTENTIAL SURPRISE: [possible content] [reason] [confidence level]"
- Include 3-5 possible surprises
</prediction>
"""
</prediction>"""
messages = [
{
"role": "user",
"content": f"Please analyze this conversation and provide a prediction following the format above:\n{chat_history}",
}
]
# Add existing user representation if available
if user_representation != "None":
messages.append(
{
"role": "user",
"content": f"Consider this existing user representation for context, but focus on current state:\n{user_representation}",
}
)
langfuse_context.update_current_observation(
input=messages, model=ANTHROPIC_MODEL
)
message = anthropic.messages.create(
model=ANTHROPIC_MODEL,
max_tokens=1000,
temperature=0,
messages=messages,
system=system_prompt,
)
print(f"tom_inference in single_prompt.py: {message.content[0].text=}")
message = message.content[0].text
return message
@ai_track("User Representation")
@observe(as_type="generation")
async def get_user_representation_single_prompt(
chat_history: str,
session_id: str,
user_representation: str = "None",
tom_inference: str = "None",
**kwargs,
) -> str:
with sentry_sdk.start_transaction(
op="user-representation-inference", name="User Representation"
):
system_prompt = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
USER_REPRESENTATION_SYSTEM_PROMPT = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
Your job is to update the existing user representation (if provided) with the new information from the conversation history and theory of mind analysis.
@ -152,36 +98,101 @@ UPDATES:
- New Information: Recent observations
- Changes: Modified interpretations
- Removals: Information no longer supported
</representation>
"""
</representation>"""
messages = []
print(f"in single_prompt.py: chat_history: {chat_history}")
print(f"in single_prompt.py: user_representation: {user_representation}")
@ai_track("Tom Inference")
@observe()
async def get_tom_inference_single_prompt(
chat_history: str,
session_id: str,
user_representation: Optional[str] = None,
**kwargs,
) -> str:
with sentry_sdk.start_transaction(op="tom-inference", name="ToM Inference"):
# Create a new model client
client = ModelClient(provider=DEF_PROVIDER, model=DEF_MODEL)
# Prepare the messages
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": f"Please analyze this conversation and provide a prediction following the format above:\n{chat_history}",
}
]
# Add existing user representation if available
if user_representation:
messages.append(
{
"role": "user",
"content": f"Consider this existing user representation for context, but focus on current state:\n{user_representation}",
}
)
langfuse_context.update_current_observation(input=messages, model=DEF_MODEL)
# Generate the response with caching enabled
try:
response = await client.generate(
messages=messages,
system=TOM_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True, # Enable caching for the system prompt
)
except Exception as e:
sentry_sdk.capture_exception(e)
logger.error(f"Error generating Tom inference: {e}")
raise e
return response
@ai_track("User Representation")
@observe()
async def get_user_representation_single_prompt(
chat_history: str,
session_id: str,
user_representation: Optional[str] = None,
tom_inference: Optional[str] = None,
**kwargs,
) -> str:
with sentry_sdk.start_transaction(
op="user-representation-inference", name="User Representation"
):
# Create a new model client
client = ModelClient(provider=DEF_PROVIDER, model=DEF_MODEL)
# Build the context message
context_str = f"CONVERSATION:\n{chat_history}\n\n"
if tom_inference != "None":
if tom_inference:
context_str += f"PREDICTION OF USER MENTAL STATE - MIGHT BE INCORRECT:\n{tom_inference}\n\n"
if user_representation != "None":
if user_representation:
context_str += f"EXISTING USER REPRESENTATION - INCOMPLETE, TO BE UPDATED:\n{user_representation}"
messages.append(
# Prepare the messages
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": f"Please analyze this information and provide an updated user representation:\n{context_str}",
}
)
]
langfuse_context.update_current_observation(
input=messages, model=ANTHROPIC_MODEL
)
message = anthropic.messages.create(
model=ANTHROPIC_MODEL,
max_tokens=1000,
temperature=0,
messages=messages,
system=system_prompt,
)
message = message.content[0].text
return message
langfuse_context.update_current_observation(input=messages, model=DEF_MODEL)
# Generate the response with caching enabled
try:
response = await client.generate(
messages=messages,
system=USER_REPRESENTATION_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True, # Enable caching for the system prompt
)
except Exception as e:
sentry_sdk.capture_exception(e)
logger.error(f"Error generating user representation: {e}")
raise e
return response

18
src/exceptions.py
View File

@ -1,14 +1,16 @@
"""
Custom exceptions for the Honcho application.
"""
from typing import Optional
class HonchoException(Exception):
"""Base exception for all Honcho-specific errors."""
status_code = 500
detail = "An unexpected error occurred"
def __init__(self, detail: Optional[str] = None, status_code: Optional[int] = None):
self.detail = detail or self.detail
self.status_code = status_code or self.status_code
@ -17,29 +19,41 @@ class HonchoException(Exception):
class ResourceNotFoundException(HonchoException):
"""Exception raised when a requested resource is not found."""
status_code = 404
detail = "Resource not found"
class ValidationException(HonchoException):
"""Exception raised when validation fails."""
status_code = 422
detail = "Validation error"
class ConflictException(HonchoException):
"""Exception raised when there's a resource conflict."""
status_code = 409
detail = "Resource conflict"
class AuthenticationException(HonchoException):
"""Exception raised when authentication fails."""
status_code = 401
detail = "Authentication failed"
class AuthorizationException(HonchoException):
"""Exception raised when authorization fails."""
status_code = 403
detail = "Not authorized to access this resource"
detail = "Not authorized to access this resource"
class DisabledException(HonchoException):
"""Exception raised when a feature is disabled."""
status_code = 405
detail = "Feature is disabled"

57
src/main.py
View File

@ -10,33 +10,65 @@ from fastapi_pagination import add_pagination
from sentry_sdk.integrations.fastapi import FastApiIntegration
from sentry_sdk.integrations.starlette import StarletteIntegration
from src.db import engine
from src.exceptions import HonchoException
from src.routers import (
apps,
collections,
documents,
keys,
messages,
metamessages,
sessions,
users,
)
from src.security import create_admin_jwt
def get_log_level(env_var="LOG_LEVEL", default="INFO"):
"""
Convert log level string from environment variable to logging module constant.
Args:
env_var: Name of the environment variable to check
default: Default log level if environment variable is not set
Returns:
int: The logging level constant (e.g., logging.INFO)
"""
log_level_str = os.getenv(env_var, default).upper()
log_levels = {
"CRITICAL": logging.CRITICAL, # 50
"ERROR": logging.ERROR, # 40
"WARNING": logging.WARNING, # 30
"INFO": logging.INFO, # 20
"DEBUG": logging.DEBUG, # 10
"NOTSET": logging.NOTSET, # 0
}
return log_levels.get(log_level_str, logging.INFO)
from .db import engine, scaffold_db
# Configure logging
logging.basicConfig(
level=logging.INFO,
level=get_log_level(),
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
# Sentry Setup
# JWT Setup
async def setup_admin_jwt():
token = create_admin_jwt()
print(f"\n ADMIN JWT: {token}\n")
# Sentry Setup
SENTRY_ENABLED = os.getenv("SENTRY_ENABLED", "False").lower() == "true"
if SENTRY_ENABLED:
sentry_sdk.init(
dsn=os.getenv("SENTRY_DSN"),
enable_tracing=True,
traces_sample_rate=0.4,
profiles_sample_rate=0.4,
integrations=[
@ -52,7 +84,6 @@ if SENTRY_ENABLED:
@asynccontextmanager
async def lifespan(app: FastAPI):
scaffold_db() # Scaffold Database on Startup
yield
await engine.dispose()
@ -60,17 +91,17 @@ async def lifespan(app: FastAPI):
app = FastAPI(
lifespan=lifespan,
servers=[
{"url": "http://127.0.0.1:8000", "description": "Local Development Server"},
{"url": "https:/demo.honcho.dev", "description": "Demo Server"},
{"url": "http://localhost:8000", "description": "Local Development Server"},
{"url": "https://demo.honcho.dev", "description": "Demo Server"},
{"url": "https://api.honcho.dev", "description": "Production SaaS Platform"},
],
title="Honcho API",
summary="An API for adding personalization to AI Apps",
description="""This API is used to store data and get insights about users for AI
applications""",
version="0.0.16",
summary="The Identity Layer for the Agentic World",
description="""Honcho is a platform for giving agents user-centric memory and social cognition""",
version="1.0.0",
contact={
"name": "Plastic Labs",
"url": "https://plasticlabs.ai",
"url": "https://honcho.dev",
"email": "hello@plasticlabs.ai",
},
license_info={
@ -99,9 +130,9 @@ app.include_router(users.router, prefix="/v1")
app.include_router(sessions.router, prefix="/v1")
app.include_router(messages.router, prefix="/v1")
app.include_router(metamessages.router, prefix="/v1")
app.include_router(metamessages.router_user_level, prefix="/v1")
app.include_router(collections.router, prefix="/v1")
app.include_router(documents.router, prefix="/v1")
app.include_router(keys.router, prefix="/v1")
# Global exception handlers

59
src/models.py
View File

@ -10,7 +10,9 @@ from sqlalchemy import (
DateTime,
ForeignKey,
Identity,
Index,
UniqueConstraint,
text,
)
from sqlalchemy.dialects.postgresql import JSONB, TEXT
from sqlalchemy.orm import Mapped, mapped_column, relationship
@ -60,12 +62,14 @@ class User(Base):
app = relationship("App", back_populates="users")
sessions = relationship("Session", back_populates="user")
collections = relationship("Collection", back_populates="user")
metamessages = relationship("Metamessage", back_populates="user")
__table_args__ = (
UniqueConstraint("name", "app_id", name="unique_name_app_user"),
CheckConstraint("length(public_id) = 21", name="public_id_length"),
CheckConstraint("length(name) <= 512", name="name_length"),
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
Index("idx_users_app_lookup", "app_id", "public_id"),
)
def __repr__(self) -> str:
@ -86,12 +90,14 @@ class Session(Base):
DateTime(timezone=True), index=True, default=func.now()
)
messages = relationship("Message", back_populates="session")
metamessages = relationship("Metamessage", back_populates="session")
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
user = relationship("User", back_populates="sessions")
__table_args__ = (
CheckConstraint("length(public_id) = 21", name="public_id_length"),
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
Index("idx_sessions_user_lookup", "user_id", "public_id"),
)
def __repr__(self) -> str:
@ -123,6 +129,12 @@ class Message(Base):
CheckConstraint("length(public_id) = 21", name="public_id_length"),
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
CheckConstraint("length(content) <= 65535", name="content_length"),
Index(
"idx_messages_session_lookup",
"session_id",
"id",
postgresql_include=["public_id", "is_user", "created_at"],
),
)
def __repr__(self) -> str:
@ -139,11 +151,21 @@ class Metamessage(Base):
)
metamessage_type: Mapped[str] = mapped_column(TEXT, index=True)
content: Mapped[str] = mapped_column(TEXT)
message_id: Mapped[str] = mapped_column(
ForeignKey("messages.public_id"), index=True
# Foreign keys - message_id is now optional
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
session_id: Mapped[str | None] = mapped_column(
ForeignKey("sessions.public_id"), index=True, nullable=True
)
message_id: Mapped[str | None] = mapped_column(
ForeignKey("messages.public_id"), index=True, nullable=True
)
# Relationships
user = relationship("User", back_populates="metamessages")
session = relationship("Session", back_populates="metamessages")
message = relationship("Message", back_populates="metamessages")
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), index=True, default=func.now()
)
@ -156,10 +178,41 @@ class Metamessage(Base):
CheckConstraint(
"length(metamessage_type) <= 512", name="metamessage_type_length"
),
# Added constraints to ensure consistency
CheckConstraint(
"(message_id IS NULL) OR (session_id IS NOT NULL)",
name="message_requires_session",
),
# Keep existing index
Index(
"idx_metamessages_lookup",
"metamessage_type",
text("id DESC"),
postgresql_include=["public_id", "message_id", "created_at"],
),
# Indices for user, session, and message lookups
Index(
"idx_metamessages_user_lookup",
"user_id",
"metamessage_type",
text("id DESC"),
),
Index(
"idx_metamessages_session_lookup",
"session_id",
"metamessage_type",
text("id DESC"),
),
Index(
"idx_metamessages_message_lookup",
"message_id",
"metamessage_type",
text("id DESC"),
),
)
def __repr__(self) -> str:
return f"Metamessages(id={self.id}, message_id={self.message_id}, metamessage_type={self.metamessage_type}, content={self.content[10:]})"
return f"Metamessages(id={self.id}, user_id={self.user_id}, session_id={self.session_id}, message_id={self.message_id}, metamessage_type={self.metamessage_type})"
class Collection(Base):

108
src/routers/apps.py
View File

@ -1,46 +1,112 @@
import logging
from typing import Optional
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import crud, schemas
from src.dependencies import db
from src.exceptions import ResourceNotFoundException
from src.security import auth
from src.exceptions import AuthenticationException, ResourceNotFoundException
from src.security import JWTParams, require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps",
tags=["apps"],
dependencies=[Depends(auth)],
)
@router.get("/{app_id}", response_model=schemas.App)
async def get_app(app_id: str, db=db):
"""Get an App by ID"""
# ResourceNotFoundException will be caught by global handler if app not found
app = await crud.get_app(db, app_id=app_id)
return app
@router.get("", response_model=schemas.App)
async def get_app(
app_id: Optional[str] = Query(
None, description="App ID to retrieve. If not provided, uses JWT token"
),
jwt_params: JWTParams = Depends(require_auth()),
db=db,
):
"""
Get an App by ID.
If app_id is provided as a query parameter, it uses that (must match JWT app_id).
Otherwise, it uses the app_id from the JWT token.
"""
# If app_id provided in query, check if it matches jwt or user is admin
if app_id:
if not jwt_params.ad and jwt_params.ap != app_id:
raise AuthenticationException("Unauthorized access to resource")
target_app_id = app_id
else:
# Use app_id from JWT
if not jwt_params.ap:
raise AuthenticationException("App ID not found in query parameter or JWT")
target_app_id = jwt_params.ap
return await crud.get_app(db, app_id=target_app_id)
@router.get("/name/{name}", response_model=schemas.App)
async def get_app_by_name(name: str, db=db):
@router.post(
"/list",
response_model=Page[schemas.App],
dependencies=[Depends(require_auth(admin=True))],
)
async def get_all_apps(
options: schemas.AppGet = Body(
..., description="Filtering and pagination options for the apps list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get all Apps"""
return await paginate(
db,
await crud.get_all_apps(
db,
reverse=reverse,
filter=options.filter,
),
)
@router.get(
"/name/{name}",
response_model=schemas.App,
dependencies=[Depends(require_auth(admin=True))],
)
async def get_app_by_name(
name: str = Path(..., description="Name of the app to retrieve"),
db=db,
):
"""Get an App by Name"""
# ResourceNotFoundException will be caught by global handler if app not found
app = await crud.get_app_by_name(db, name=name)
return app
@router.post("", response_model=schemas.App)
async def create_app(app: schemas.AppCreate, db=db):
@router.post(
"", response_model=schemas.App, dependencies=[Depends(require_auth(admin=True))]
)
async def create_app(
app: schemas.AppCreate = Body(..., description="App creation parameters"),
db=db,
):
"""Create a new App"""
honcho_app = await crud.create_app(db, app=app)
return honcho_app
@router.get("/get_or_create/{name}", response_model=schemas.App)
async def get_or_create_app(name: str, db=db):
@router.get(
"/get_or_create/{name}",
response_model=schemas.App,
dependencies=[Depends(require_auth(admin=True))],
)
async def get_or_create_app(
name: str = Path(..., description="Name of the app to get or create"),
db=db,
):
"""Get or Create an App"""
try:
app = await crud.get_app_by_name(db=db, name=name)
@ -51,10 +117,14 @@ async def get_or_create_app(name: str, db=db):
return app
@router.put("/{app_id}", response_model=schemas.App)
@router.put(
"/{app_id}",
response_model=schemas.App,
dependencies=[Depends(require_auth(app_id="app_id"))],
)
async def update_app(
app_id: str,
app: schemas.AppUpdate,
app_id: str = Path(..., description="ID of the app to update"),
app: schemas.AppUpdate = Body(..., description="Updated app parameters"),
db=db,
):
"""Update an App"""

151
src/routers/collections.py
View File

@ -1,26 +1,82 @@
from typing import Optional
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import crud, schemas
from src.dependencies import db
from src.security import auth
from src.exceptions import AuthenticationException
from src.security import JWTParams, require_auth
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/collections",
tags=["collections"],
dependencies=[Depends(auth)],
)
@router.post("/list", response_model=Page[schemas.Collection])
@router.get(
"",
response_model=schemas.Collection,
)
async def get_collection(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: Optional[str] = Query(
None, description="Collection ID to retrieve. If not provided, uses JWT token"
),
jwt_params: JWTParams = Depends(require_auth()),
db=db,
):
"""
Get a specific collection for a user.
If collection_id is provided as a query parameter, it uses that (must match JWT collection_id).
Otherwise, it uses the collection_id from the JWT token.
"""
# Verify JWT has access to the requested resource
if not jwt_params.ad:
if jwt_params.ap is not None and jwt_params.ap != app_id:
raise AuthenticationException("Unauthorized access to resource")
if jwt_params.us is not None and jwt_params.us != user_id:
raise AuthenticationException("Unauthorized access to resource")
# If collection_id provided in query, check if it matches jwt or user is admin
if collection_id:
if (
not jwt_params.ad
and jwt_params.co is not None
and jwt_params.co != collection_id
):
raise AuthenticationException("Unauthorized access to resource")
target_collection_id = collection_id
else:
# Use collection_id from JWT
if not jwt_params.co:
raise AuthenticationException(
"Collection ID not found in query parameter or JWT"
)
target_collection_id = jwt_params.co
# Let crud function handle the ResourceNotFoundException
return await crud.get_collection_by_id(
db, app_id=app_id, collection_id=target_collection_id, user_id=user_id
)
@router.post(
"/list",
response_model=Page[schemas.Collection],
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def get_collections(
app_id: str,
user_id: str,
options: schemas.CollectionGet,
reverse: Optional[bool] = False,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
options: schemas.CollectionGet = Body(
..., description="Filtering options for the collections list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get All Collections for a User"""
@ -32,11 +88,15 @@ async def get_collections(
)
@router.get("/name/{name}", response_model=schemas.Collection)
@router.get(
"/name/{name}",
response_model=schemas.Collection,
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def get_collection_by_name(
app_id: str,
user_id: str,
name: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
name: str = Path(..., description="Name of the collection to retrieve"),
db=db,
) -> schemas.Collection:
"""Get a Collection by Name"""
@ -46,25 +106,17 @@ async def get_collection_by_name(
return honcho_collection
@router.get("/{collection_id}", response_model=schemas.Collection)
async def get_collection_by_id(
app_id: str,
user_id: str,
collection_id: str,
db=db,
) -> schemas.Collection:
"""Get a Collection by ID"""
honcho_collection = await crud.get_collection_by_id(
db, app_id=app_id, user_id=user_id, collection_id=collection_id
)
return honcho_collection
@router.post("", response_model=schemas.Collection)
@router.post(
"",
response_model=schemas.Collection,
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def create_collection(
app_id: str,
user_id: str,
collection: schemas.CollectionCreate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection: schemas.CollectionCreate = Body(
..., description="Collection creation parameters"
),
db=db,
):
"""Create a new Collection"""
@ -75,12 +127,24 @@ async def create_collection(
)
@router.put("/{collection_id}", response_model=schemas.Collection)
@router.put(
"/{collection_id}",
response_model=schemas.Collection,
dependencies=[
Depends(
require_auth(
app_id="app_id", user_id="user_id", collection_id="collection_id"
)
)
],
)
async def update_collection(
app_id: str,
user_id: str,
collection_id: str,
collection: schemas.CollectionUpdate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection to update"),
collection: schemas.CollectionUpdate = Body(
..., description="Updated collection parameters"
),
db=db,
):
"Update a Collection's name or metadata"
@ -96,11 +160,20 @@ async def update_collection(
return honcho_collection
@router.delete("/{collection_id}")
@router.delete(
"/{collection_id}",
dependencies=[
Depends(
require_auth(
app_id="app_id", user_id="user_id", collection_id="collection_id"
)
)
],
)
async def delete_collection(
app_id: str,
user_id: str,
collection_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection to delete"),
db=db,
):
"""Delete a Collection and its documents"""

112
src/routers/documents.py
View File

@ -1,33 +1,42 @@
import json
import logging
from collections.abc import Sequence
from typing import Optional
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import crud, schemas
from src.dependencies import db
from src.exceptions import ResourceNotFoundException, ValidationException
from src.security import auth
from src.security import require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/collections/{collection_id}/documents",
tags=["documents"],
dependencies=[Depends(auth)],
dependencies=[
Depends(
require_auth(
app_id="app_id", user_id="user_id", collection_id="collection_id"
)
)
],
)
@router.post("/list", response_model=Page[schemas.Document])
async def get_documents(
app_id: str,
user_id: str,
collection_id: str,
options: schemas.DocumentGet,
reverse: Optional[bool] = False,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
options: schemas.DocumentGet = Body(
..., description="Filtering options for the documents list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get all of the Documents in a Collection"""
@ -40,22 +49,23 @@ async def get_documents(
filter=options.filter,
reverse=reverse,
)
return await paginate(db, documents_query)
except ValueError as e:
logger.warning(f"Failed to get documents for collection {collection_id}: {str(e)}")
raise ResourceNotFoundException("Collection not found or does not belong to user") from e
logger.warning(
f"Failed to get documents for collection {collection_id}: {str(e)}"
)
raise ResourceNotFoundException(
"Collection not found or does not belong to user"
) from e
@router.get(
"/{document_id}",
response_model=schemas.Document,
)
@router.get("/{document_id}", response_model=schemas.Document)
async def get_document(
app_id: str,
user_id: str,
collection_id: str,
document_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
document_id: str = Path(..., description="ID of the document to retrieve"),
db=db,
):
"""Get a document by ID"""
@ -71,10 +81,12 @@ async def get_document(
@router.post("/query", response_model=Sequence[schemas.Document])
async def query_documents(
app_id: str,
user_id: str,
collection_id: str,
options: schemas.DocumentQuery,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
options: schemas.DocumentQuery = Body(
..., description="Query parameters for document search"
),
db=db,
):
"""Cosine Similarity Search for Documents"""
@ -84,7 +96,7 @@ async def query_documents(
filter = options.filter
if options.filter == {}:
filter = None
documents = await crud.query_documents(
db=db,
app_id=app_id,
@ -94,20 +106,24 @@ async def query_documents(
filter=filter,
top_k=top_k,
)
logger.info(f"Query documents successful for collection {collection_id}")
return documents
except ValueError as e:
logger.error(f"Error querying documents in collection {collection_id}: {str(e)}")
logger.error(
f"Error querying documents in collection {collection_id}: {str(e)}"
)
raise ValidationException("Error querying documents") from e
@router.post("", response_model=schemas.Document)
async def create_document(
app_id: str,
user_id: str,
collection_id: str,
document: schemas.DocumentCreate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
document: schemas.DocumentCreate = Body(
..., description="Document creation parameters"
),
db=db,
):
"""Embed text as a vector and create a Document"""
@ -122,8 +138,12 @@ async def create_document(
logger.info(f"Document created successfully in collection {collection_id}")
return document_obj
except ValueError as e:
logger.warning(f"Failed to create document in collection {collection_id}: {str(e)}")
raise ResourceNotFoundException("Collection not found or does not belong to user") from e
logger.warning(
f"Failed to create document in collection {collection_id}: {str(e)}"
)
raise ResourceNotFoundException(
"Collection not found or does not belong to user"
) from e
@router.put(
@ -131,18 +151,22 @@ async def create_document(
response_model=schemas.Document,
)
async def update_document(
app_id: str,
user_id: str,
collection_id: str,
document_id: str,
document: schemas.DocumentUpdate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
document_id: str = Path(..., description="ID of the document to update"),
document: schemas.DocumentUpdate = Body(
..., description="Updated document parameters"
),
db=db,
):
"""Update the content and/or the metadata of a Document"""
if document.content is None and document.metadata is None:
logger.warning(f"Document update attempted with empty content and metadata for document {document_id}")
logger.warning(
f"Document update attempted with empty content and metadata for document {document_id}"
)
raise ValidationException("Content and metadata cannot both be None")
try:
updated_document = await crud.update_document(
db,
@ -161,10 +185,10 @@ async def update_document(
@router.delete("/{document_id}")
async def delete_document(
app_id: str,
user_id: str,
collection_id: str,
document_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
document_id: str = Path(..., description="ID of the document to delete"),
db=db,
):
"""Delete a Document by ID"""

58
src/routers/keys.py Normal file
View File

@ -0,0 +1,58 @@
import datetime
import logging
import os
from fastapi import APIRouter, Depends, Query
from src.exceptions import DisabledException, ValidationException
from src.security import (
JWTParams,
create_jwt,
require_auth,
)
logger = logging.getLogger(__name__)
USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true"
router = APIRouter(
prefix="/keys",
tags=["keys"],
dependencies=[Depends(require_auth(admin=True))],
)
@router.post("")
async def create_key(
app_id: str | None = Query(None, description="ID of the app to scope the key to"),
user_id: str | None = Query(None, description="ID of the user to scope the key to"),
session_id: str | None = Query(
None, description="ID of the session to scope the key to"
),
collection_id: str | None = Query(
None, description="ID of the collection to scope the key to"
),
expires_at: datetime.datetime | None = None,
):
"""Create a new Key"""
if not USE_AUTH:
raise DisabledException()
# Validate that at least one parameter is provided for proper scoping
if not any([app_id, user_id, session_id, collection_id]):
raise ValidationException(
"At least one of app_id, user_id, session_id, or collection_id must be provided"
)
key_str = create_jwt(
JWTParams(
exp=expires_at.isoformat() if expires_at else None,
ap=app_id,
us=user_id,
se=session_id,
co=collection_id,
)
)
return {
"key": key_str,
}

76
src/routers/messages.py
View File

@ -2,7 +2,7 @@ import logging
import os
from typing import List, Optional
from fastapi import APIRouter, BackgroundTasks, Depends
from fastapi import APIRouter, BackgroundTasks, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from sqlalchemy.sql import insert
@ -12,14 +12,18 @@ from src.db import SessionLocal
from src.dependencies import db
from src.exceptions import ResourceNotFoundException
from src.models import QueueItem
from src.security import auth
from src.security import require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/sessions/{session_id}/messages",
tags=["messages"],
dependencies=[Depends(auth)],
dependencies=[
Depends(
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
)
],
)
@ -47,6 +51,11 @@ async def enqueue(payload: dict | list[dict]):
user_id=payload[0]["user_id"],
session_id=payload[0]["session_id"],
)
if not session:
logger.warning(
f"Session {payload[0]['session_id']} not found, skipping enqueue"
)
return
except ResourceNotFoundException:
logger.warning(
f"Session {payload[0]['session_id']} not found, skipping enqueue"
@ -97,6 +106,11 @@ async def enqueue(payload: dict | list[dict]):
user_id=payload["user_id"],
session_id=payload["session_id"],
)
if not session:
logger.warning(
f"Session {payload['session_id']} not found, skipping enqueue"
)
return
except ResourceNotFoundException:
logger.warning(
f"Session {payload['session_id']} not found, skipping enqueue"
@ -137,11 +151,13 @@ async def enqueue(payload: dict | list[dict]):
@router.post("", response_model=schemas.Message)
async def create_message_for_session(
app_id: str,
user_id: str,
session_id: str,
message: schemas.MessageCreate,
background_tasks: BackgroundTasks,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
message: schemas.MessageCreate = Body(
..., description="Message creation parameters"
),
db=db,
):
"""Adds a message to a session"""
@ -175,11 +191,13 @@ async def create_message_for_session(
@router.post("/batch", response_model=List[schemas.Message])
async def create_batch_messages_for_session(
app_id: str,
user_id: str,
session_id: str,
batch: schemas.MessageBatchCreate,
background_tasks: BackgroundTasks,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
batch: schemas.MessageBatchCreate = Body(
..., description="Batch of messages to create"
),
db=db,
):
"""Bulk create messages for a session while maintaining order. Maximum 100 messages per batch."""
@ -222,11 +240,15 @@ async def create_batch_messages_for_session(
@router.post("/list", response_model=Page[schemas.Message])
async def get_messages(
app_id: str,
user_id: str,
session_id: str,
options: schemas.MessageGet,
reverse: Optional[bool] = False,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
options: schemas.MessageGet = Body(
..., description="Filtering options for the messages list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get all messages for a session"""
@ -252,10 +274,10 @@ async def get_messages(
@router.get("/{message_id}", response_model=schemas.Message)
async def get_message(
app_id: str,
user_id: str,
session_id: str,
message_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
message_id: str = Path(..., description="ID of the message to retrieve"),
db=db,
):
"""Get a Message by ID"""
@ -270,11 +292,13 @@ async def get_message(
@router.put("/{message_id}", response_model=schemas.Message)
async def update_message(
app_id: str,
user_id: str,
session_id: str,
message_id: str,
message: schemas.MessageUpdate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
message_id: str = Path(..., description="ID of the message to update"),
message: schemas.MessageUpdate = Body(
..., description="Updated message parameters"
),
db=db,
):
"""Update the metadata of a Message"""
@ -291,4 +315,4 @@ async def update_message(
return updated_message
except ValueError as e:
logger.warning(f"Failed to update message {message_id}: {str(e)}")
raise ResourceNotFoundException("Message or session not found") from e
raise ResourceNotFoundException("Message not found") from e

132
src/routers/metamessages.py
View File

@ -1,105 +1,87 @@
import logging
from typing import Optional
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import crud, schemas
from src.dependencies import db
from src.exceptions import ResourceNotFoundException, ValidationException
from src.security import auth
from src.security import require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/sessions/{session_id}/metamessages",
tags=["metamessages"],
dependencies=[Depends(auth)],
)
router_user_level = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/metamessages",
tags=["metamessages"],
dependencies=[Depends(auth)],
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
@router.post("", response_model=schemas.Metamessage)
async def create_metamessage(
app_id: str,
user_id: str,
session_id: str,
metamessage: schemas.MetamessageCreate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
metamessage: schemas.MetamessageCreate = Body(
..., description="Metamessage creation parameters"
),
db=db,
):
"""Adds a message to a session"""
"""
Create a new metamessage associated with a user.
Optionally link to a session and message by providing those IDs in the request body.
"""
try:
metamessage_obj = await crud.create_metamessage(
db,
user_id=user_id,
metamessage=metamessage,
app_id=app_id,
user_id=user_id,
session_id=session_id,
)
logger.info(f"Metamessage created successfully for session {session_id}")
logger.info(f"Metamessage created successfully for user {user_id}")
return metamessage_obj
except ValueError as e:
logger.warning(f"Failed to create metamessage for session {session_id}: {str(e)}")
raise ResourceNotFoundException("Session not found") from e
except (ResourceNotFoundException, ValidationException) as e:
logger.warning(f"Failed to create metamessage: {str(e)}")
raise
@router.post("/list", response_model=Page[schemas.Metamessage])
async def get_metamessages(
app_id: str,
user_id: str,
session_id: str,
options: schemas.MetamessageGet,
reverse: Optional[bool] = False,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
options: schemas.MetamessageGet = Body(
..., description="Filtering options for the metamessages list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get all messages for a session"""
"""
Get metamessages with flexible filtering.
- Filter by user only: No additional parameters needed
- Filter by session: Provide session_id
- Filter by message: Provide message_id (and session_id)
- Filter by type: Provide metamessage_type
- Filter by metadata: Provide filter object
"""
try:
metamessages_query = await crud.get_metamessages(
db,
app_id=app_id,
user_id=user_id,
session_id=session_id,
session_id=options.session_id,
message_id=options.message_id,
metamessage_type=options.metamessage_type,
filter=options.filter,
reverse=reverse,
)
return await paginate(db, metamessages_query)
except ValueError as e:
logger.warning(f"Failed to get metamessages for session {session_id}: {str(e)}")
raise ResourceNotFoundException("Session not found") from e
@router_user_level.post("/list", response_model=Page[schemas.Metamessage])
async def get_metamessages_by_user(
app_id: str,
user_id: str,
options: schemas.MetamessageGetUserLevel,
reverse: Optional[bool] = False,
db=db,
):
"""Paginate through the user metamessages for a user"""
try:
metamessages_query = await crud.get_metamessages(
db,
app_id=app_id,
user_id=user_id,
metamessage_type=options.metamessage_type,
reverse=reverse,
filter=options.filter,
)
return await paginate(db, metamessages_query)
except ValueError as e:
logger.warning(f"Failed to get metamessages for user {user_id}: {str(e)}")
raise ResourceNotFoundException("User not found") from e
except (ResourceNotFoundException, ValidationException) as e:
logger.warning(f"Failed to get metamessages: {str(e)}")
raise
@router.get(
@ -107,25 +89,23 @@ async def get_metamessages_by_user(
response_model=schemas.Metamessage,
)
async def get_metamessage(
app_id: str,
user_id: str,
session_id: str,
message_id: str,
metamessage_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
metamessage_id: str = Path(..., description="ID of the metamessage to retrieve"),
db=db,
):
"""Get a specific Metamessage by ID"""
honcho_metamessage = await crud.get_metamessage(
db,
app_id=app_id,
session_id=session_id,
user_id=user_id,
message_id=message_id,
metamessage_id=metamessage_id,
)
if honcho_metamessage is None:
logger.warning(f"Metamessage {metamessage_id} not found for message {message_id}")
raise ResourceNotFoundException(f"Metamessage with ID {metamessage_id} not found")
logger.warning(f"Metamessage {metamessage_id} not found")
raise ResourceNotFoundException(
f"Metamessage with ID {metamessage_id} not found"
)
return honcho_metamessage
@ -134,29 +114,25 @@ async def get_metamessage(
response_model=schemas.Metamessage,
)
async def update_metamessage(
app_id: str,
user_id: str,
session_id: str,
metamessage_id: str,
metamessage: schemas.MetamessageUpdate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
metamessage_id: str = Path(..., description="ID of the metamessage to update"),
metamessage: schemas.MetamessageUpdate = Body(
..., description="Updated metamessage parameters"
),
db=db,
):
"""Update's the metadata of a metamessage"""
if metamessage.metadata is None:
logger.warning(f"Update attempted with empty metadata for metamessage {metamessage_id}")
raise ValidationException("Metamessage metadata cannot be empty")
"""Update a metamessage's metadata, type, or relationships"""
try:
updated_metamessage = await crud.update_metamessage(
db,
metamessage=metamessage,
app_id=app_id,
user_id=user_id,
session_id=session_id,
metamessage_id=metamessage_id,
)
logger.info(f"Metamessage {metamessage_id} updated successfully")
return updated_metamessage
except ValueError as e:
except (ResourceNotFoundException, ValidationException) as e:
logger.warning(f"Failed to update metamessage {metamessage_id}: {str(e)}")
raise ResourceNotFoundException("Session or metamessage not found") from e
raise

271
src/routers/sessions.py
View File

@ -1,32 +1,92 @@
import logging
from typing import Optional
from anthropic import MessageStreamManager
from fastapi import APIRouter, Depends
from anthropic import AsyncMessageStreamManager
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi.exceptions import HTTPException
from fastapi.responses import StreamingResponse
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import agent, crud, schemas
from src.dependencies import db
from src.exceptions import ResourceNotFoundException, ValidationException
from src.security import auth
from src.exceptions import (
AuthenticationException,
ResourceNotFoundException,
ValidationException,
)
from src.security import JWTParams, require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/sessions",
tags=["sessions"],
dependencies=[Depends(auth)],
)
@router.post("/list", response_model=Page[schemas.Session])
@router.get(
"",
response_model=schemas.Session,
)
async def get_session(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: Optional[str] = Query(
None, description="Session ID to retrieve. If not provided, uses JWT token"
),
jwt_params: JWTParams = Depends(require_auth()),
db=db,
):
"""
Get a specific session for a user.
If session_id is provided as a query parameter, it uses that (must match JWT session_id).
Otherwise, it uses the session_id from the JWT token.
"""
# Verify JWT has access to the requested resource
if not jwt_params.ad:
if jwt_params.ap is not None and jwt_params.ap != app_id:
raise AuthenticationException("Unauthorized access to resource")
if jwt_params.us is not None and jwt_params.us != user_id:
raise AuthenticationException("Unauthorized access to resource")
# If session_id provided in query, check if it matches jwt or user is admin
if session_id:
if (
not jwt_params.ad
and jwt_params.se is not None
and jwt_params.se != session_id
):
raise AuthenticationException("Unauthorized access to resource")
target_session_id = session_id
else:
# Use session_id from JWT
if not jwt_params.se:
raise AuthenticationException(
"Session ID not found in query parameter or JWT"
)
target_session_id = jwt_params.se
# Let crud function handle the ResourceNotFoundException
return await crud.get_session(
db, app_id=app_id, session_id=target_session_id, user_id=user_id
)
@router.post(
"/list",
response_model=Page[schemas.Session],
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def get_sessions(
app_id: str,
user_id: str,
options: schemas.SessionGet,
reverse: Optional[bool] = False,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
options: schemas.SessionGet = Body(
..., description="Filtering and pagination options for the sessions list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get All Sessions for a User"""
@ -43,11 +103,17 @@ async def get_sessions(
)
@router.post("", response_model=schemas.Session)
@router.post(
"",
response_model=schemas.Session,
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def create_session(
app_id: str,
user_id: str,
session: schemas.SessionCreate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session: schemas.SessionCreate = Body(
..., description="Session creation parameters"
),
db=db,
):
"""Create a Session for a User"""
@ -62,12 +128,22 @@ async def create_session(
raise ValidationException(str(e)) from e
@router.put("/{session_id}", response_model=schemas.Session)
@router.put(
"/{session_id}",
response_model=schemas.Session,
dependencies=[
Depends(
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
)
],
)
async def update_session(
app_id: str,
user_id: str,
session_id: str,
session: schemas.SessionUpdate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session to update"),
session: schemas.SessionUpdate = Body(
..., description="Updated session parameters"
),
db=db,
):
"""Update the metadata of a Session"""
@ -82,11 +158,18 @@ async def update_session(
raise ResourceNotFoundException("Session not found") from e
@router.delete("/{session_id}")
@router.delete(
"/{session_id}",
dependencies=[
Depends(
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
)
],
)
async def delete_session(
app_id: str,
user_id: str,
session_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session to delete"),
db=db,
):
"""Delete a session by marking it as inactive"""
@ -101,88 +184,92 @@ async def delete_session(
raise ResourceNotFoundException("Session not found") from e
@router.get("/{session_id}", response_model=schemas.Session)
async def get_session(
app_id: str,
user_id: str,
session_id: str,
db=db,
):
"""Get a specific session for a user by ID"""
honcho_session = await crud.get_session(
db, app_id=app_id, session_id=session_id, user_id=user_id
)
if honcho_session is None:
logger.warning(f"Session {session_id} not found for user {user_id}")
raise ResourceNotFoundException(f"Session with ID {session_id} not found")
return honcho_session
@router.post("/{session_id}/chat", response_model=schemas.AgentChat)
async def chat(
app_id: str,
user_id: str,
session_id: str,
query: schemas.AgentQuery,
):
"""Chat with the Dialectic API"""
return await agent.chat(
app_id=app_id, user_id=user_id, session_id=session_id, query=query
)
@router.post(
"/{session_id}/chat/stream",
"/{session_id}/chat",
response_model=schemas.DialecticResponse,
responses={
200: {
"description": "Chat stream",
"content": {
"text/event-stream": {"schema": {"type": "string", "format": "binary"}}
},
}
"description": "Response to a question informed by Honcho's User Representation",
"content": {"text/event-stream": {}},
},
},
dependencies=[
Depends(
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
)
],
)
async def get_chat_stream(
app_id: str,
user_id: str,
session_id: str,
query: schemas.AgentQuery,
async def chat(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
options: schemas.DialecticOptions = Body(
..., description="Dialectic Endpoint Parameters"
),
):
"""Stream Results from the Dialectic API"""
async def parse_stream():
stream = await agent.chat(
"""Chat with the Dialectic API"""
if not options.stream:
return await agent.chat(
app_id=app_id,
user_id=user_id,
session_id=session_id,
query=query,
stream=True,
queries=options.queries,
)
if type(stream) is MessageStreamManager:
with stream as stream_manager:
for text in stream_manager.text_stream:
yield text
else:
return StreamingResponse(
content=parse_stream(), media_type="text/event-stream", status_code=200
)
async def parse_stream():
try:
stream = await agent.chat(
app_id=app_id,
user_id=user_id,
session_id=session_id,
queries=options.queries,
stream=True,
)
if type(stream) is AsyncMessageStreamManager:
async with stream as stream_manager:
async for text in stream_manager.text_stream:
yield text
except Exception as e:
logger.error(f"Error in stream: {str(e)}")
raise HTTPException(status_code=500, detail=str(e)) from e
return StreamingResponse(
content=parse_stream(), media_type="text/event-stream", status_code=200
)
@router.get("/{session_id}/clone", response_model=schemas.Session)
@router.get(
"/{session_id}/clone",
response_model=schemas.Session,
dependencies=[
Depends(
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
)
],
)
async def clone_session(
app_id: str,
user_id: str,
session_id: str,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session to clone"),
db=db,
message_id: Optional[str] = None,
deep_copy: bool = False,
message_id: Optional[str] = Query(
None, description="Message ID to cut off the clone at"
),
deep_copy: bool = Query(False, description="Whether to deep copy metamessages"),
):
"""Clone a session for a user, optionally will deep clone metamessages as well"""
return await crud.clone_session(
db,
app_id=app_id,
user_id=user_id,
original_session_id=session_id,
cutoff_message_id=message_id,
deep_copy=deep_copy,
)
"""Clone a session, optionally up to a specific message"""
try:
cloned_session = await crud.clone_session(
db,
app_id=app_id,
user_id=user_id,
original_session_id=session_id,
cutoff_message_id=message_id,
deep_copy=deep_copy,
)
logger.info(f"Session {session_id} cloned successfully")
return cloned_session
except ValueError as e:
logger.warning(f"Failed to clone session {session_id}: {str(e)}")
raise ResourceNotFoundException("Session not found") from e

119
src/routers/users.py
View File

@ -1,29 +1,34 @@
import logging
from typing import Optional
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import crud, schemas
from src.dependencies import db
from src.exceptions import (
AuthenticationException,
ResourceNotFoundException,
)
from src.security import auth
from src.security import JWTParams, require_auth
logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users",
tags=["users"],
dependencies=[Depends(auth)],
)
@router.post("", response_model=schemas.User)
@router.post(
"",
response_model=schemas.User,
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def create_user(
app_id: str,
user: schemas.UserCreate,
app_id: str = Path(..., description="ID of the app"),
user: schemas.UserCreate = Body(..., description="User creation parameters"),
db=db,
):
"""Create a new User"""
@ -31,11 +36,17 @@ async def create_user(
return user_obj
@router.post("/list", response_model=Page[schemas.User])
@router.post(
"/list",
response_model=Page[schemas.User],
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def get_users(
app_id: str,
options: schemas.UserGet,
reverse: bool = False,
app_id: str = Path(..., description="ID of the app"),
options: schemas.UserGet = Body(
..., description="Filtering options for the users list"
),
reverse: bool = Query(False, description="Whether to reverse the order of results"),
db=db,
):
"""Get All Users for an App"""
@ -45,10 +56,55 @@ async def get_users(
)
@router.get("/name/{name}", response_model=schemas.User)
@router.get(
"",
response_model=schemas.User,
)
async def get_user(
app_id: str = Path(..., description="ID of the app"),
user_id: Optional[str] = Query(
None, description="User ID to retrieve. If not provided, users JWT token"
),
jwt_params: JWTParams = Depends(require_auth()),
db=db,
):
"""
Get a User by ID
If user_id is provided as a query parameter, it uses that (must match JWT app_id).
Otherwise, it uses the user_id from the JWT token.
"""
# validate app query param
if not jwt_params.ad and jwt_params.ap is not None and jwt_params.ap != app_id:
raise AuthenticationException("Unauthorized access to resource")
if user_id:
if not jwt_params.ad and jwt_params.us is not None and jwt_params.us != user_id:
raise AuthenticationException("Unauthorized access to resource")
target_user_id = user_id
else:
# Use user_id from JWT
if not jwt_params.us:
raise AuthenticationException("User ID not found in query parameter or JWT")
target_user_id = jwt_params.us
user = await crud.get_user(db, app_id=app_id, user_id=target_user_id)
return user
@router.get(
"/name/{name}",
response_model=schemas.User,
dependencies=[
Depends(
require_auth(
app_id="app_id",
)
)
],
)
async def get_user_by_name(
app_id: str,
name: str,
app_id: str = Path(..., description="ID of the app"),
name: str = Path(..., description="Name of the user to retrieve"),
db=db,
):
"""Get a User by name"""
@ -56,19 +112,22 @@ async def get_user_by_name(
return user
@router.get("/{user_id}", response_model=schemas.User)
async def get_user(
app_id: str,
user_id: str,
@router.get(
"/get_or_create/{name}",
response_model=schemas.User,
dependencies=[
Depends(
require_auth(
app_id="app_id",
)
)
],
)
async def get_or_create_user(
app_id: str = Path(..., description="ID of the app"),
name: str = Path(..., description="Name of the user to get or create"),
db=db,
):
"""Get a User by ID"""
user = await crud.get_user(db, app_id=app_id, user_id=user_id)
return user
@router.get("/get_or_create/{name}", response_model=schemas.User)
async def get_or_create_user(app_id: str, name: str, db=db):
"""Get a User or create a new one by the input name"""
try:
user = await crud.get_user_by_name(db, app_id=app_id, name=name)
@ -81,11 +140,15 @@ async def get_or_create_user(app_id: str, name: str, db=db):
return user
@router.put("/{user_id}", response_model=schemas.User)
@router.put(
"/{user_id}",
response_model=schemas.User,
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
async def update_user(
app_id: str,
user_id: str,
user: schemas.UserUpdate,
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user to update"),
user: schemas.UserUpdate = Body(..., description="Updated user parameters"),
db=db,
):
"""Update a User's name and/or metadata"""

35
src/schemas.py
View File

@ -13,6 +13,10 @@ class AppCreate(AppBase):
metadata: dict = {}
class AppGet(AppBase):
filter: dict | None = None
class AppUpdate(AppBase):
name: str | None = None
metadata: dict | None = None
@ -172,23 +176,21 @@ class MetamessageBase(BaseModel):
class MetamessageCreate(MetamessageBase):
metamessage_type: Annotated[str, Field(min_length=1, max_length=50)]
content: Annotated[str, Field(min_length=0, max_length=50000)]
message_id: str
session_id: str | None = None
message_id: str | None = None
metadata: dict = {}
class MetamessageGet(MetamessageBase):
metamessage_type: str | None = None
session_id: str | None = None
message_id: str | None = None
filter: dict | None = None
class MetamessageGetUserLevel(MessageBase):
filter: dict | None = None
metamessage_type: str | None = None
class MetamessageUpdate(MetamessageBase):
message_id: str
session_id: str | None = None
message_id: str | None = None
metamessage_type: str | None = None
metadata: dict | None = None
@ -198,7 +200,9 @@ class Metamessage(MetamessageBase):
id: str
metamessage_type: str
content: str
message_id: str
user_id: str
session_id: str | None
message_id: str | None
h_metadata: dict = Field(exclude=True)
metadata: dict
created_at: datetime.datetime
@ -317,27 +321,30 @@ class Document(DocumentBase):
)
class AgentQuery(BaseModel):
class DialecticOptions(BaseModel):
queries: str | list[str]
stream: bool = False
@field_validator('queries')
@field_validator("queries")
def validate_queries(cls, v):
MAX_STRING_LENGTH = 10000
MAX_LIST_LENGTH = 25
if isinstance(v, str):
if len(v) > MAX_STRING_LENGTH:
raise ValueError('Query too long')
raise ValueError("Query too long")
elif isinstance(v, list):
if len(v) > MAX_LIST_LENGTH:
raise ValueError('Too many queries')
raise ValueError("Too many queries")
if any(len(q) > MAX_STRING_LENGTH for q in v):
raise ValueError('One or more queries too long')
raise ValueError("One or more queries too long")
return v
class AgentChat(BaseModel):
class DialecticResponse(BaseModel):
content: str
class MessageBatchCreate(BaseModel):
"""Schema for batch message creation with a max of 100 messages"""
messages: list[MessageCreate] = Field(..., max_length=100)

208
src/security.py
View File

@ -1,28 +1,214 @@
import datetime
import logging
import os
from typing import Annotated
from typing import Annotated, Optional
from fastapi import Depends
import jwt
from fastapi import Depends, Request
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession
from src.dependencies import get_db
from .exceptions import AuthenticationException
logger = logging.getLogger(__name__)
USE_AUTH_SERVICE = os.getenv("USE_AUTH_SERVICE", "False").lower() == "true"
SECRET_KEY = os.getenv("SECRET_KEY", "test")
USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true"
AUTH_JWT_SECRET = os.getenv("AUTH_JWT_SECRET", "") if USE_AUTH else ""
if USE_AUTH and AUTH_JWT_SECRET == "":
print(
"\n ERROR: No JWT secret provided. Set the AUTH_JWT_SECRET environment variable.\n"
)
exit(1)
security = HTTPBearer(
auto_error=False,
)
#
# jwt params
# all optional, used to produce tokens valid for different routes
# hierarchy: app > user > ( session / collection )
# routes that involve a 'name' parameter require permissions for the parent object
# name routes are considered 'queries' as names are mutable properties
#
# note: add routes without parameters that assume the most immediately scoped key is providing
#
class JWTParams(BaseModel):
"""
JWT parameters used to produce tokens valid for different routes.
Hierarchy: app > user > (session / collection)
All routers require at least the most tightly scoped parameter.
Routes will accept a JWT with a scope higher in the hierarchy.
Names shortened to minimize token size. Timestamp is included
so that many unique tokens can be generated for the same resource.
Note that the timestamp itself is not used for security, and can
be omitted, such as when Honcho generates the initial admin JWT.
Fields (all optional other than `t`):
`t`: a string timestamp of when the JWT was created
`exp`: a string timestamp of when the JWT expires (optional)
`ad`: a boolean flag indicating if the JWT is an admin JWT
`ap`: (string) app id
`us`: (string) user id
`se`: (string) session id
`co`: (string) collection id
"""
t: str = datetime.datetime.now().isoformat()
exp: Optional[str] = None
ad: Optional[bool] = None
ap: Optional[str] = None
us: Optional[str] = None
se: Optional[str] = None
co: Optional[str] = None
def create_admin_jwt() -> str:
"""Create a JWT for admin operations."""
params = JWTParams(t="", ad=True)
key = create_jwt(params)
return key
def create_jwt(params: JWTParams) -> str:
"""Create a JWT token from the given parameters."""
payload = {k: v for k, v in params.__dict__.items() if v is not None}
return jwt.encode(payload, AUTH_JWT_SECRET.encode("utf-8"), algorithm="HS256")
async def verify_jwt(token: str) -> JWTParams:
"""Verify a JWT token and return the decoded parameters."""
params = JWTParams()
try:
decoded = jwt.decode(
token, AUTH_JWT_SECRET.encode("utf-8"), algorithms=["HS256"]
)
if "t" in decoded:
params.t = decoded["t"]
if "exp" in decoded:
params.exp = decoded["exp"]
if (
params.exp
and datetime.datetime.fromisoformat(params.exp)
< datetime.datetime.now()
):
raise AuthenticationException("JWT expired")
if "ad" in decoded:
params.ad = decoded["ad"]
if "ap" in decoded:
params.ap = decoded["ap"]
if "us" in decoded:
params.us = decoded["us"]
if "se" in decoded:
params.se = decoded["se"]
if "co" in decoded:
params.co = decoded["co"]
return params
except jwt.PyJWTError:
raise AuthenticationException("Invalid JWT") from None
def require_auth(
admin: Optional[bool] = None,
app_id: Optional[str] = None,
user_id: Optional[str] = None,
session_id: Optional[str] = None,
collection_id: Optional[str] = None,
):
"""
Generate a dependency that requires authentication for the given parameters.
"""
async def auth_dependency(
request: Request,
credentials: HTTPAuthorizationCredentials = Depends(security),
db: AsyncSession = Depends(get_db),
):
app_id_param = (
request.path_params.get(app_id) or request.query_params.get(app_id)
if app_id
else None
)
user_id_param = (
request.path_params.get(user_id) or request.query_params.get(user_id)
if user_id
else None
)
session_id_param = (
request.path_params.get(session_id) or request.query_params.get(session_id)
if session_id
else None
)
collection_id_param = (
request.path_params.get(collection_id)
or request.query_params.get(collection_id)
if collection_id
else None
)
return await auth(
credentials=credentials,
admin=admin,
app_id=app_id_param,
user_id=user_id_param,
session_id=session_id_param,
collection_id=collection_id_param,
)
return auth_dependency
async def auth(
credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)],
):
if not USE_AUTH_SERVICE:
return True
if not credentials or credentials.credentials != SECRET_KEY:
logger.warning("Invalid access token attempt")
raise AuthenticationException("Invalid access token")
return {"message": "OK"}
admin: Optional[bool] = None,
app_id: Optional[str] = None,
user_id: Optional[str] = None,
session_id: Optional[str] = None,
collection_id: Optional[str] = None,
) -> JWTParams:
"""Authenticate the given JWT and return the decoded parameters."""
if not USE_AUTH:
return JWTParams(t="", ad=True)
if not credentials or not credentials.credentials:
logger.warning("No access token provided")
raise AuthenticationException("No access token provided")
jwt_params = await verify_jwt(credentials.credentials)
# based on api operation, verify api key based on that key's permissions
if jwt_params.ad:
return jwt_params
if admin:
raise AuthenticationException("Resource requires admin privileges")
# Check if the JWT has direct access to the requested resource
# For session or collection level access
if session_id and jwt_params.se == session_id:
return jwt_params
if collection_id and jwt_params.co == collection_id:
return jwt_params
# For user level access - can access all sessions/collections under this user
if user_id and jwt_params.us == user_id:
return jwt_params
# For app level access - can access all users/sessions/collections under this app
if app_id and jwt_params.ap == app_id:
return jwt_params
if any([session_id, collection_id, user_id, app_id]):
print([session_id, collection_id, user_id, app_id])
print(jwt_params)
raise AuthenticationException("JWT not permissioned for this resource")
# Route did not specify any parameters, so it should parse parameters itself
return jwt_params

21
src/utils/__init__.py Normal file
View File

@ -0,0 +1,21 @@
"""
Utility modules for the Honcho app.
"""
import re
def parse_xml_content(text: str, tag: str) -> str:
"""
Extract content from XML-like tags in a string.
Args:
text: The text containing XML-like tags
tag: The tag name to extract content from
Returns:
The content between the opening and closing tags, or an empty string if not found
"""
pattern = f"<{tag}>(.*?)</{tag}>"
match = re.search(pattern, text, re.DOTALL)
return match.group(1).strip() if match else ""

392
src/utils/history.py Normal file
View File

@ -0,0 +1,392 @@
import logging
from enum import Enum
from typing import Optional, Union, cast
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src.utils.model_client import ModelClient, ModelProvider
from .. import models
logger = logging.getLogger(__name__)
# Export the public functions
__all__ = [
"get_session_summaries",
"get_messages_since_message",
"get_messages_since_latest_summary",
"create_summary",
"save_summary_metamessage",
"get_summarized_history",
"should_create_summary",
"MESSAGES_PER_SHORT_SUMMARY",
"MESSAGES_PER_LONG_SUMMARY",
"SummaryType",
]
# Configuration constants for summaries
MESSAGES_PER_SHORT_SUMMARY = 20 # How often to create short summaries
MESSAGES_PER_LONG_SUMMARY = 60 # How often to create long summaries
# The types of metamessages to use for summaries
class SummaryType(Enum):
SHORT = "honcho_chat_summary_short"
LONG = "honcho_chat_summary_long"
# Default model settings for summary generation
DEFAULT_PROVIDER = ModelProvider.GEMINI
DEFAULT_MODEL = "gemini-2.0-flash-lite"
async def get_session_summaries(
db: AsyncSession,
session_id: str,
summary_type: SummaryType = SummaryType.SHORT,
only_latest: bool = False,
) -> Union[list[models.Metamessage], Optional[models.Metamessage]]:
"""
Get summaries for a given session.
Args:
db: Database session
session_id: The session ID
summary_type: Type of summary to retrieve ("short" or "long")
only_latest: Whether to return only the latest summary
Returns:
If only_latest is True: The most recent summary metamessage, or None if none exists
If only_latest is False: A list of all summary metamessages for the session
"""
# Determine the metamessage type based on summary_type
metamessage_type = (
SummaryType.SHORT.value
if summary_type == SummaryType.SHORT
else SummaryType.LONG.value
)
stmt = (
select(models.Metamessage)
.where(models.Metamessage.session_id == session_id)
.where(models.Metamessage.metamessage_type == metamessage_type)
.order_by(models.Metamessage.id.desc())
)
if only_latest:
stmt = stmt.limit(1)
result = await db.execute(stmt)
# Always return a metamessage instance or None
return result.scalar_one_or_none()
else:
result = await db.execute(stmt)
return list(result.scalars().all())
async def get_messages_since_message(
db: AsyncSession, session_id: str, message_id: Optional[str] = None
) -> list[models.Message]:
"""
Get all messages since a specific message.
Args:
db: Database session
session_id: The session ID
message_id: The reference message ID
Returns:
List of messages after the reference message or all messages if message_id is None
"""
# Base query for messages in this session
query = (
select(models.Message)
.where(models.Message.session_id == session_id)
.order_by(models.Message.id)
)
# If we have a reference message ID, filter to get newer messages
if message_id:
# First, get the ID of the message
message_id_subquery = (
select(models.Message.id)
.where(models.Message.public_id == message_id)
.scalar_subquery()
)
# Then filter to only get messages with higher IDs
query = query.where(models.Message.id > message_id_subquery)
# Execute query
result = await db.execute(query)
return list(result.scalars().all())
async def create_summary(
messages: list[models.Message],
previous_summary: Optional[str] = None,
summary_type: SummaryType = SummaryType.SHORT,
) -> str:
"""
Generate a summary of the provided messages using an LLM.
Args:
messages: List of messages to summarize
previous_summary: Optional previous summary to provide context
summary_type: Type of summary to create ("short" or "long")
Returns:
A summary of the conversation
"""
# Combine messages into a conversation format
conversation = "\n".join(
[
f"{'human' if msg.is_user else 'assistant'}: {msg.content}"
for msg in messages
]
)
# Adjust system prompt based on summary type
if summary_type == SummaryType.LONG:
system_prompt = """You are a system that creates comprehensive summaries of conversations.
Focus on capturing:
1. Key facts and information shared
2. User preferences, opinions, and questions
3. Important context and requests
4. Core topics discussed in detail
5. User's apparent emotional state and personality traits
6. Important themes and patterns across the conversation
It is very important that you clearly distinguish between the user's messages and the assistant's messages, and that only the user's literal words are attributed to them.
Provide a thorough and detailed summary that captures the essence of the conversation.
Your summary should serve as a comprehensive record of the important information in this conversation.
Return only the summary without any explanation or meta-commentary."""
else: # short summary
system_prompt = """You are a system that summarizes parts of a conversation to create a concise and accurate summary.
Focus on capturing:
1. Key facts and information shared
2. User preferences, opinions, and questions
3. Important context and requests
4. Core topics discussed
5. User's apparent emotional state
It is very important that you clearly distinguish between the user's messages and the assistant's messages, and that only the user's literal words are attributed to them.
Provide a concise, factual summary that captures the essence of the conversation.
Your summary should be detailed enough to serve as context for future messages,
but brief enough to be helpful.
Return only the summary without any explanation or meta-commentary."""
# Include previous summary if available
if previous_summary:
user_prompt = f"""Here is a previous summary of the conversation:
{previous_summary}
Now please summarize these additional messages, incorporating the context from the previous summary.
Your summary should summarize the entire conversation in a self-contained way, such that someone could read it and understand the entire conversation.
{conversation}
Provide a {"comprehensive" if summary_type == SummaryType.LONG else "concise"} summary that captures both the previous context and the new information."""
else:
user_prompt = f"""Please summarize the following conversation segment:
{conversation}
Provide a {"comprehensive" if summary_type == SummaryType.LONG else "concise"} summary that captures the key points and context."""
# Create a model client
client = ModelClient(provider=DEFAULT_PROVIDER, model=DEFAULT_MODEL)
# Generate the summary
llm_messages = [{"role": "user", "content": user_prompt}]
try:
summary = await client.generate(
messages=llm_messages,
system=system_prompt,
max_tokens=1000
if summary_type == SummaryType.SHORT
else 2000, # Allow longer responses for long summaries
temperature=0.0,
use_caching=True,
)
return summary
except Exception as e:
logger.error(f"Error generating summary: {str(e)}")
# Fallback to a basic summary in case of error
return f"Conversation with {len(messages)} messages about {messages[-1].content[:30]}..."
async def save_summary_metamessage(
db: AsyncSession,
user_id: str,
session_id: str,
message_id: str,
summary_content: str,
message_count: int,
summary_type: SummaryType = SummaryType.SHORT,
) -> models.Metamessage:
"""
Save a summary as a metamessage.
Args:
db: Database session
user_id: User ID
session_id: Session ID
message_id: The ID of the most recent message being summarized
summary_content: The summary text to save
message_count: Number of messages covered by this summary
summary_type: Type of summary to save
Returns:
The created metamessage
"""
# Get the metamessage_type value from the enum
metamessage_type = summary_type.value
# Create and save the metamessage
metamessage = models.Metamessage(
user_id=user_id,
session_id=session_id,
message_id=message_id,
metamessage_type=metamessage_type,
content=summary_content,
h_metadata={"message_count": message_count, "summary_type": summary_type.name},
)
db.add(metamessage)
await db.commit()
logger.info(
f"Saved {summary_type.name.lower()} summary metamessage for session {session_id} covering {message_count} messages"
)
return metamessage
async def get_full_history(
db: AsyncSession, session_id: str
) -> tuple[str, list[models.Message]]:
"""
Get all messages for a given session.
"""
messages = await get_messages_since_message(db, session_id)
return format_messages(messages), messages
async def get_summarized_history(
db: AsyncSession, session_id: str, summary_type: SummaryType = SummaryType.SHORT
) -> tuple[str, list[models.Message], Optional[models.Metamessage]]:
"""
Get a summarized version of the chat history by combining the latest summary
with all messages since that summary.
Args:
db: Database session
session_id: The session ID
summary_type: Type of summary to get ("short" or "long")
Returns:
A tuple containing:
- String formatted history text with summary and recent messages
- List of messages since the latest summary
- The latest summary metamessage, or None if no summary exists
"""
# Get messages since the latest summary and the summary itself
messages, latest_summary = await get_messages_since_latest_summary(
db, session_id, summary_type
)
# Format messages
messages_text = format_messages(messages)
# We know latest_summary is either a Metamessage or None because of the type
# narrowing in get_messages_since_latest_summary
if latest_summary:
# Combine summary with recent messages
history_text = f"[CONVERSATION SUMMARY: {latest_summary.content}]\n\n[RECENT MESSAGES]\n{messages_text}"
else:
# No summary available, return just the messages
history_text = messages_text
return history_text, messages, latest_summary
async def get_messages_since_latest_summary(
db: AsyncSession, session_id: str, summary_type: SummaryType = SummaryType.SHORT
) -> tuple[list[models.Message], Optional[models.Metamessage]]:
"""
Get all messages since the latest summary for a session.
This is a convenience method that combines:
1. Getting the latest summary for the session
2. Getting all messages since that summary
Args:
db: Database session
session_id: The session ID
summary_type: Type of summary to get ("short" or "long")
Returns:
A tuple containing:
- List of messages since the latest summary (or all messages if no summary exists)
- The latest summary metamessage, or None if no summary exists
"""
# Get the latest summary (will be a Metamessage instance or None)
# We include only_latest=True to get a single Metamessage (not a list)
summary = await get_session_summaries(
db, session_id, summary_type=summary_type, only_latest=True
)
# Type narrowing - summary is now either None or a Metamessage instance
latest_summary = cast(Optional[models.Metamessage], summary)
# Check if we have a valid summary with a message_id
if latest_summary is not None:
messages = await get_messages_since_message(
db, session_id, latest_summary.message_id
)
return messages, latest_summary
else:
messages = await get_messages_since_message(db, session_id)
return messages, None
async def should_create_summary(
db: AsyncSession, session_id: str, summary_type: SummaryType = SummaryType.SHORT
) -> tuple[bool, list[models.Message], Optional[models.Metamessage]]:
"""
Determine if a new summary should be created for this session.
Args:
db: Database session
session_id: The session ID
summary_type: Type of summary to check for ("short" or "long")
Returns:
Tuple containing:
- Boolean indicating whether a summary should be created
- List of messages to be included in the summary
- The latest summary of the requested type, or None if no summary exists
"""
messages, latest_summary = await get_messages_since_latest_summary(
db, session_id, summary_type
)
threshold = (
MESSAGES_PER_SHORT_SUMMARY
if summary_type == SummaryType.SHORT
else MESSAGES_PER_LONG_SUMMARY
)
should_create = len(messages) >= threshold
return should_create, messages, latest_summary
def format_messages(messages: list[models.Message]) -> str:
"""
Format a list of messages into a string.
"""
if len(messages) == 0:
return ""
return "\n".join(
[f"{'user' if msg.is_user else 'assistant'}: {msg.content}" for msg in messages]
)

541
src/utils/model_client.py Normal file
View File

@ -0,0 +1,541 @@
"""
Utility functions for interacting with various language model APIs.
"""
import os
from enum import Enum
from typing import Any, Optional, Protocol
import sentry_sdk
from anthropic import AsyncAnthropic
from dotenv import load_dotenv
from google import genai
from google.genai import types as genai_types
from langfuse.decorators import langfuse_context, observe
# from openai import AsyncOpenAI
from langfuse.openai import AsyncOpenAI
# Load environment variables
load_dotenv()
# Supported model providers
class ModelProvider(str, Enum):
ANTHROPIC = "anthropic"
OPENAI = "openai"
OPENROUTER = "openrouter"
CEREBRAS = "cerebras"
GROQ = "groq"
GEMINI = "gemini"
# Add other providers as needed
# Default models for each provider
DEFAULT_MODELS = {
ModelProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
ModelProvider.OPENAI: "gpt-4o",
ModelProvider.OPENROUTER: "meta-llama/Llama-3.3-70B-Instruct",
ModelProvider.CEREBRAS: "llama-3.3-70b",
ModelProvider.GROQ: "llama-3.3-70b-versatile",
ModelProvider.GEMINI: "gemini-2.0-flash-lite",
}
OPENAI_COMPATIBLE_PROVIDERS = [
ModelProvider.OPENAI,
ModelProvider.OPENROUTER,
ModelProvider.CEREBRAS,
ModelProvider.GROQ,
]
DEFAULT_TEMPERATURE = 0.0
DEFAULT_MAX_TOKENS = 1000
class Message(Protocol):
"""Protocol for a message that works with any provider."""
role: str
content: str
class ModelClient:
"""A client for interacting with various language model APIs."""
def __init__(
self,
provider: ModelProvider = ModelProvider.ANTHROPIC,
model: Optional[str] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
):
"""
Initialize the model client.
Args:
provider: The model provider to use
model: The specific model to use, or None to use the default model for the provider
api_key: The API key to use, or None to read from environment variables
base_url: Custom base URL for the API endpoints (used for OpenRouter)
"""
self.provider = provider
self.model = model or DEFAULT_MODELS[provider]
self.base_url = base_url
self.openai_client = None
self.gemini_client = None
# Setup provider-specific clients
if provider == ModelProvider.ANTHROPIC:
self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
if not self.api_key:
raise ValueError("Anthropic API key is required")
self.client = AsyncAnthropic(api_key=self.api_key)
elif provider in OPENAI_COMPATIBLE_PROVIDERS:
self.api_key = api_key or os.getenv("OPENAI_COMPATIBLE_API_KEY")
self.base_url = base_url or os.getenv("OPENAI_COMPATIBLE_BASE_URL")
if not self.api_key:
raise ValueError("OpenAI-compatible API key is required")
self.openai_client = AsyncOpenAI(
api_key=self.api_key, base_url=self.base_url
)
elif provider == ModelProvider.GEMINI:
self.api_key = api_key or os.getenv("GEMINI_API_KEY")
if not self.api_key:
raise ValueError("Gemini API key is required")
self.gemini_client = genai.Client(api_key=self.api_key)
else:
raise ValueError(f"Unsupported provider: {provider}")
def create_message(self, role: str, content: str) -> dict[str, str]:
"""
Create a message that works with the current provider.
Args:
role: The role of the message (e.g., "user", "assistant")
content: The message content
Returns:
A message compatible with the current provider
"""
# For now, just return a dictionary that works with both Anthropic and OpenAI
return {"role": role, "content": content}
@observe()
async def generate(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> str:
"""
Generate a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
max_tokens: Maximum number of tokens to generate
temperature: Temperature for generation
extra_headers: Optional headers to add to the request
use_caching: Whether to use provider-side caching for the response
Returns:
The generated text
"""
with sentry_sdk.start_transaction(
op="llm-api", name=f"{self.provider} API Call"
):
# Log to langfuse
langfuse_context.update_current_observation(
input=messages, model=self.model
)
if self.provider == ModelProvider.ANTHROPIC:
return await self._generate_anthropic(
messages,
system,
max_tokens,
temperature,
extra_headers,
use_caching,
)
elif self.provider in OPENAI_COMPATIBLE_PROVIDERS:
return await self._generate_openai(
messages, system, max_tokens, temperature
)
elif self.provider == ModelProvider.GEMINI:
return await self._generate_gemini(
messages, system, max_tokens, temperature
)
else:
raise ValueError(f"Unsupported provider: {self.provider}")
@observe(as_type="generation")
async def _generate_anthropic(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> str:
"""Generate a response using the Anthropic API."""
if not self.client:
raise ValueError("Anthropic client not initialized.")
params = {
"model": self.model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
params["system"] = [
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"},
}
]
else:
params["system"] = system
langfuse_context.update_current_observation(input=messages, model=self.model)
response = await self.client.messages.create(**params)
# Extract the text from the response
if response.content and len(response.content) > 0:
content_block = response.content[0]
# Check content_block by checking for attribute 'type' instead of using isinstance
if (
content_block
and hasattr(content_block, "type")
and content_block.type == "text"
):
return content_block.text
return str(content_block)
return ""
@observe(as_type="generation")
async def _generate_openai(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> str:
"""Generate text using OpenAI or OpenRouter API."""
if not self.openai_client:
raise ValueError("OpenAI client not initialized")
# Prepare messages
formatted_messages = []
# Add system message if provided
if system:
formatted_messages.append({"role": "system", "content": system})
# Add the rest of the messages
formatted_messages.extend(messages)
# Make the API call with the OpenAI client
response = await self.openai_client.chat.completions.create(
model=self.model,
messages=formatted_messages,
max_tokens=max_tokens,
temperature=temperature,
)
# Extract the generated text
choice = response.choices[0]
if choice and choice.message and choice.message.content:
return choice.message.content
return ""
@observe()
async def stream(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> Any:
"""
Stream a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
max_tokens: Maximum number of tokens to generate
temperature: Temperature for generation
extra_headers: Optional headers to add to the request
use_caching: Whether to use provider-side caching for the response
Returns:
A streaming response from the provider
"""
with sentry_sdk.start_transaction(
op="llm-api-stream", name=f"{self.provider} API Stream"
):
# Log to langfuse
langfuse_context.update_current_observation(
input=messages, model=self.model
)
if self.provider == ModelProvider.ANTHROPIC:
return await self._stream_anthropic(
messages,
system,
max_tokens,
temperature,
extra_headers,
use_caching,
)
elif self.provider in OPENAI_COMPATIBLE_PROVIDERS:
return await self._stream_openai(
messages, system, max_tokens, temperature
)
elif self.provider == ModelProvider.GEMINI:
return await self._stream_gemini(
messages, system, max_tokens, temperature
)
else:
raise ValueError(f"Unsupported provider: {self.provider}")
@observe(as_type="generation")
async def _stream_anthropic(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> Any:
"""Stream text using Anthropic API."""
if not self.client:
raise ValueError("Anthropic client not initialized.")
params = {
"model": self.model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
params["system"] = [
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"},
}
]
else:
params["system"] = system
langfuse_context.update_current_observation(input=messages, model=self.model)
# Return the stream directly without awaiting it
return self.client.messages.stream(**params)
@observe(as_type="generation")
async def _stream_openai(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> Any:
"""Stream text using OpenAI or OpenRouter API."""
if not self.openai_client:
raise ValueError("OpenAI client not initialized")
# Prepare messages
formatted_messages = []
# Add system message if provided
if system:
formatted_messages.append({"role": "system", "content": system})
# Add the rest of the messages
formatted_messages.extend(messages)
# Make the API call with the OpenAI client
stream = await self.openai_client.chat.completions.create(
model=self.model,
messages=formatted_messages,
max_tokens=max_tokens,
temperature=temperature,
stream=True,
)
return stream
@observe(as_type="generation")
async def _generate_gemini(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> str:
"""Generate text using Gemini API."""
if not self.gemini_client:
raise ValueError("Gemini client not initialized")
# Format messages for Gemini
gemini_messages = []
# Convert messages to Gemini format
for message in messages:
role = message["role"]
# Map roles to what Gemini expects
if role == "user":
gemini_role = "user"
elif role == "assistant":
gemini_role = "model"
else:
# Skip system messages as they're handled through config
continue
gemini_messages.append(
genai_types.Content(
role=gemini_role,
parts=[genai_types.Part.from_text(text=message["content"])],
)
)
# Set generation config
generate_content_config = genai_types.GenerateContentConfig(
temperature=temperature,
max_output_tokens=max_tokens,
response_mime_type="text/plain",
)
# Add system instruction if provided
if system:
generate_content_config.system_instruction = system
# Make the API call
if not gemini_messages:
# If we have no messages but have a system prompt, create a default user message
if system:
default_content = genai_types.Content(
role="user",
parts=[
genai_types.Part.from_text(
text="Please respond based on the system instructions."
)
],
)
# model = self.gemini_client.get_model(self.model)
response = await self.gemini_client.aio.models.generate_content(
model=self.model,
contents=default_content,
config=generate_content_config,
)
else:
raise ValueError("No messages provided for Gemini generation")
else:
# Normal case with messages
# model = get_model(self.model)
response = await self.gemini_client.aio.models.generate_content(
model=self.model,
contents=gemini_messages
if len(gemini_messages) > 1
else gemini_messages[0],
config=generate_content_config,
)
# Extract text from response
if response and response.text:
return response.text
return ""
@observe(as_type="generation")
async def _stream_gemini(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> Any:
"""Stream text using Gemini API."""
if not self.gemini_client:
raise ValueError("Gemini client not initialized")
# Format messages for Gemini
gemini_messages = []
# Convert messages to Gemini format
for message in messages:
role = message["role"]
# Map roles to what Gemini expects
if role == "user":
gemini_role = "user"
elif role == "assistant":
gemini_role = "model"
else:
# Skip system messages as they're handled through config
continue
gemini_messages.append(
genai_types.Content(
role=gemini_role,
parts=[genai_types.Part.from_text(text=message["content"])],
)
)
# Set generation config
generate_content_config = genai_types.GenerateContentConfig(
temperature=temperature,
max_output_tokens=max_tokens,
response_mime_type="text/plain",
)
# Add system instruction if provided
if system:
generate_content_config.system_instruction = system
# Make the streaming API call
if not gemini_messages:
# If we have no messages but have a system prompt, create a default user message
if system:
default_content = genai_types.Content(
role="user",
parts=[
genai_types.Part.from_text(
text="Please respond based on the system instructions."
)
],
)
stream = await self.gemini_client.aio.models.generate_content_stream(
model=self.model,
contents=default_content,
config=generate_content_config,
)
else:
raise ValueError("No messages provided for Gemini streaming")
else:
# Normal case with messages
stream = await self.gemini_client.aio.models.generate_content_stream(
model=self.model,
contents=gemini_messages
if len(gemini_messages) > 1
else gemini_messages[0],
config=generate_content_config,
)
return stream

117
tests/conftest.py
View File

@ -1,7 +1,8 @@
import logging # noqa: I001
import os
import sys
import jwt
from nanoid import generate as generate_nanoid
from unittest.mock import patch, MagicMock, AsyncMock
import pytest
import pytest_asyncio
@ -9,7 +10,7 @@ from fastapi import Request
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine, AsyncSession
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.engine.url import make_url
from sqlalchemy.exc import OperationalError, ProgrammingError
from sqlalchemy_utils import create_database, database_exists, drop_database
@ -18,22 +19,45 @@ from src import models
from src.db import Base
from src.dependencies import get_db
from src.exceptions import HonchoException
from src.security import create_admin_jwt, create_jwt, JWTParams
from src.main import app
# Create a custom handler that doesn't get closed prematurely
class TestHandler(logging.Handler):
def __init__(self):
super().__init__()
self.records = []
def emit(self, record):
self.records.append(record)
# Setup logging with our custom handler
test_handler = TestHandler()
logging.basicConfig(
level=logging.DEBUG,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
stream=sys.stdout, # This ensures the output goes to stdout
handlers=[test_handler],
)
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
# Test database URL
# TODO use environment variable
CONNECTION_URI = make_url(os.getenv("CONNECTION_URI"))
CONNECTION_URI = make_url(
os.getenv(
"CONNECTION_URI",
"postgresql+psycopg://postgres:postgres@localhost:5432/postgres",
)
)
TEST_DB_URL = CONNECTION_URI.set(database="test_db")
DEFAULT_DB_URL = str(CONNECTION_URI.set(database="postgres"))
# Test API authorization
USE_AUTH = os.getenv("USE_AUTH", "False").lower() == "true"
AUTH_JWT_SECRET = os.getenv("AUTH_JWT_SECRET", "test-secret")
def create_test_database(db_url):
"""Helper function create a database if it does not already exist
@ -66,7 +90,7 @@ async def setup_test_database(db_url):
Returns:
engine: SQLAlchemy engine
"""
engine = create_async_engine(str(db_url))
engine = create_async_engine(str(db_url), echo=True)
async with engine.connect() as conn:
try:
logger.info("Attempting to create pgvector extension...")
@ -94,7 +118,10 @@ async def db_engine():
create_test_database(TEST_DB_URL)
engine = await setup_test_database(TEST_DB_URL)
# Drop all tables first to ensure clean state
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.drop_all)
# Then create all tables with current models
await conn.run_sync(Base.metadata.create_all)
yield engine
@ -114,7 +141,7 @@ async def db_session(db_engine):
@pytest.fixture(scope="function")
def client(db_session):
async def client(db_session):
"""Create a FastAPI TestClient for the scope of a single test function"""
# Register exception handlers for tests
@ -130,9 +157,49 @@ def client(db_session):
app.dependency_overrides[get_db] = override_get_db
with TestClient(app) as c:
if USE_AUTH:
# give the test client the admin JWT
c.headers["Authorization"] = f"Bearer {create_admin_jwt()}"
yield c
def create_invalid_jwt() -> str:
return jwt.encode({"ad": "invalid"}, "this is not the secret", algorithm="HS256")
@pytest.fixture(
params=[
("none", None), # No auth
("invalid", create_invalid_jwt), # Invalid JWT
("empty", lambda: create_jwt(JWTParams())), # Empty JWT
("admin", create_admin_jwt), # Admin JWT
]
)
def auth_client(client, request, monkeypatch):
"""
Fixture that provides a client with different authentication states.
Always ensures USE_AUTH is set to True.
"""
# Ensure USE_AUTH is always True for this fixture
import src.routers.keys as keys_module
import src.security as security
monkeypatch.setattr(keys_module, "USE_AUTH", "true")
monkeypatch.setattr(security, "USE_AUTH", "true")
# Clear any existing Authorization header
client.headers.pop("Authorization", None)
auth_type, token_func = request.param
client.auth_type = auth_type
if token_func is not None:
token = token_func()
client.headers["Authorization"] = f"Bearer {token}"
return client
@pytest_asyncio.fixture(scope="function")
async def sample_data(db_session):
"""Helper function to create test data"""
@ -149,3 +216,41 @@ async def sample_data(db_session):
yield test_app, test_user
await db_session.rollback()
@pytest.fixture(autouse=True)
def mock_langfuse():
"""Mock Langfuse decorator and context during tests"""
with (
patch("langfuse.decorators.observe") as mock_observe,
patch("langfuse.decorators.langfuse_context") as mock_context,
):
# Mock the decorator to just return the function
mock_observe.return_value = lambda func: func
# Mock the context object
mock_context_obj = MagicMock()
mock_context_obj.update_current_observation = MagicMock()
mock_context_obj.update_current_trace = MagicMock()
mock_context.return_value = mock_context_obj
# Disable httpx logging during tests
logging.getLogger("httpx").setLevel(logging.WARNING)
yield
# Clean up logging handlers
for handler in logging.getLogger().handlers[:]:
if isinstance(handler, TestHandler):
handler.close()
logging.getLogger().removeHandler(handler)
@pytest.fixture(autouse=True)
def mock_openai_embeddings():
"""Mock OpenAI embeddings API calls for testing"""
with patch("src.crud.openai_client.embeddings.create") as mock_create:
mock_response = AsyncMock()
mock_response.data = [MagicMock(embedding=[0.1] * 1536)]
mock_create.return_value = mock_response
yield mock_create

50
tests/routes/test_apps.py
View File

@ -1,3 +1,4 @@
import pytest
from nanoid import generate as generate_nanoid
@ -37,7 +38,7 @@ def test_get_or_create_app(client):
response = client.get(f"/v1/apps/name/{name}")
assert response.status_code == 404
assert "detail" in response.json()
# This should create the app
response = client.get(f"/v1/apps/get_or_create/{name}")
assert response.status_code == 200
@ -48,23 +49,23 @@ def test_get_or_create_app(client):
def test_get_or_create_existing_app(client):
name = str(generate_nanoid())
# App doesn't exist yet
response = client.get(f"/v1/apps/name/{name}")
assert response.status_code == 404
# Create the app
response = client.post(
"/v1/apps", json={"name": name, "metadata": {"key": "value"}}
)
assert response.status_code == 200
app1 = response.json()
# Now get_or_create should find the existing app
response = client.get(f"/v1/apps/get_or_create/{name}")
assert response.status_code == 200
app2 = response.json()
# Both should be the same app
assert app1["name"] == app2["name"]
assert app1["id"] == app2["id"]
@ -73,13 +74,46 @@ def test_get_or_create_existing_app(client):
def test_get_app_by_id(client, sample_data):
test_app, _ = sample_data
response = client.get(f"/v1/apps/{test_app.public_id}")
response = client.get(f"/v1/apps?app_id={test_app.public_id}")
assert response.status_code == 200
data = response.json()
assert data["name"] == test_app.name
assert data["id"] == str(test_app.public_id)
@pytest.mark.asyncio
async def test_get_all_apps(client, db_session, sample_data):
test_app, test_user = sample_data
# create a test app with metadata
response = client.post(
"/v1/apps",
json={
"name": "test_app",
"metadata": {"test_key": "test_value"},
},
)
response = client.post(
"/v1/apps/list",
json={},
)
assert response.status_code == 200
data = response.json()
assert "items" in data
assert len(data["items"]) > 0
response = client.post(
"/v1/apps/list",
json={"filter": {"test_key": "test_value"}},
)
assert response.status_code == 200
data = response.json()
assert "items" in data
assert len(data["items"]) > 0
assert data["items"][0]["metadata"] == {"test_key": "test_value"}
def test_get_app_by_name(client, sample_data):
test_app, _ = sample_data
response = client.get(f"/v1/apps/name/{test_app.name}")
@ -107,10 +141,10 @@ def test_create_duplicate_app_name(client):
name = str(generate_nanoid())
response = client.post("/v1/apps", json={"name": name})
assert response.status_code == 200
# Try to create another app with the same name
response = client.post("/v1/apps", json={"name": name})
# Should get a ConflictException with 409 status
assert response.status_code == 409
data = response.json()

57
tests/routes/test_collections.py
View File

@ -1,3 +1,4 @@
import pytest
from nanoid import generate as generate_nanoid
@ -25,7 +26,7 @@ def test_get_collection_by_id(client, sample_data) -> None:
data = response.json()
# Get the collection
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}"
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections?collection_id={data['id']}"
)
assert response.status_code == 200
data = response.json()
@ -127,6 +128,58 @@ def test_delete_collection(client, sample_data) -> None:
)
assert response.status_code == 200
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}"
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections?collection_id={data['id']}"
)
assert response.status_code == 404
@pytest.mark.asyncio
async def test_cannot_delete_honcho_collection(client, sample_data, db_session) -> None:
test_app, test_user = sample_data
# Make the protected "honcho" collection
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "honcho", "metadata": {}},
)
assert response.status_code == 422 # Should get validation error when trying to create
# Try to create it directly through API to test delete protection
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection", "metadata": {}},
)
assert response.status_code == 200
data = response.json()
collection_id = data["id"]
# Update the collection to have the reserved name
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}",
json={"name": "honcho", "metadata": {}},
)
assert response.status_code == 422 # Should get validation error
# Create the protected collection using the internal method directly with db_session
from src.crud import create_user_protected_collection
# Create the protected collection
honcho_collection = await create_user_protected_collection(
db_session, app_id=test_app.public_id, user_id=test_user.public_id
)
await db_session.flush()
# Get the protected collection
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/name/honcho"
)
assert response.status_code == 200
data = response.json()
assert data["name"] == "honcho"
protected_id = data["id"]
# Try to delete the protected collection
response = client.delete(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{protected_id}"
)
assert response.status_code == 422 # Should get validation error
assert "reserved name" in response.json()["detail"].lower()

51
tests/routes/test_documents.py
View File

@ -1,3 +1,4 @@
import pytest
from nanoid import generate as generate_nanoid
@ -10,7 +11,7 @@ def test_update_document_validation_error(client, sample_data):
)
assert response.status_code == 200
collection = response.json()
# Create a document
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents",
@ -18,13 +19,13 @@ def test_update_document_validation_error(client, sample_data):
)
assert response.status_code == 200
document = response.json()
# Try to update the document with empty content and metadata
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/{document['id']}",
json={"content": None, "metadata": None},
)
# Should get a ValidationException with 422 status
assert response.status_code == 422
data = response.json()
@ -214,3 +215,47 @@ def test_delete_document(client, sample_data):
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/{document['id']}"
)
assert response.status_code == 404
@pytest.mark.asyncio
async def test_cannot_delete_documents_in_honcho_collection(client, sample_data, db_session):
test_app, test_user = sample_data
# Create the protected collection using the internal method directly with db_session
from src.crud import create_user_protected_collection
# Create the protected honcho collection
honcho_collection = await create_user_protected_collection(
db_session, app_id=test_app.public_id, user_id=test_user.public_id
)
await db_session.flush()
# Get the protected collection
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/name/honcho"
)
assert response.status_code == 200
collection = response.json()
assert collection["name"] == "honcho"
# Create a document in the protected collection
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents",
json={"content": "protected document", "metadata": {}},
)
assert response.status_code == 200
document = response.json()
# Try to delete the document from the protected collection
response = client.delete(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/{document['id']}"
)
assert response.status_code == 422 # Should get validation error
assert "honcho" in response.json()["detail"].lower()
# The document should still exist
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/{document['id']}"
)
assert response.status_code == 200
assert response.json()["content"] == "protected document"

63
tests/routes/test_keys.py Normal file
View File

@ -0,0 +1,63 @@
def test_create_key_no_params(auth_client):
"""Test creating a key with no parameters"""
response = auth_client.post("/v1/keys")
# Only admin JWT should be allowed
if auth_client.auth_type == "admin":
# key with no params should fail
assert response.status_code == 422
else:
assert response.status_code == 401
def test_create_key_with_params(auth_client, sample_data):
"""Test creating a key with specific parameters"""
test_app, test_user = sample_data
if auth_client.auth_type != "admin":
return # Skip test if not admin authentication
# Test with app_id
response = auth_client.post("/v1/keys", params={"app_id": test_app.public_id})
assert response.status_code == 200
assert "key" in response.json()
# Test with app_id and user_id
response = auth_client.post(
"/v1/keys",
params={"app_id": test_app.public_id, "user_id": test_user.public_id},
)
assert response.status_code == 200
assert "key" in response.json()
# Test with session_id and collection_id
response = auth_client.post(
"/v1/keys",
params={
"app_id": test_app.public_id,
"user_id": test_user.public_id,
"session_id": "test-session",
"collection_id": "test-collection",
},
)
assert response.status_code == 200
assert "key" in response.json()
def test_create_key_with_expires_at(auth_client, sample_data):
"""Test creating a key with an expiration date"""
response = auth_client.post("/v1/keys", params={"expires_at": "2025-01-01"})
# Only admin JWT should be allowed
if auth_client.auth_type == "admin":
# key with no params should fail
assert response.status_code == 422
return
else:
assert response.status_code == 401
test_app, _ = sample_data
# assert that the key is expired
response = auth_client.post("/v1/keys", params={"app_id": test_app.public_id})
assert response.status_code == 401

19
tests/routes/test_messages.py
View File

@ -111,6 +111,7 @@ async def test_update_message(client, db_session, sample_data):
data = response.json()
assert data["metadata"] == {"new_key": "new_value"}
@pytest.mark.asyncio
async def test_update_message_empty_metadata(client, db_session, sample_data):
test_app, test_user = sample_data
@ -148,8 +149,9 @@ async def test_create_batch_messages(client, db_session, sample_data):
{
"content": f"Test message {i}",
"is_user": i % 2 == 0, # Alternating user/non-user messages
"metadata": {"batch_index": i}
} for i in range(3)
"metadata": {"batch_index": i},
}
for i in range(3)
]
}
@ -157,13 +159,13 @@ async def test_create_batch_messages(client, db_session, sample_data):
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/batch",
json=test_messages,
)
assert response.status_code == 200
data = response.json()
# Verify the response contains all messages
assert len(data) == 3
# Verify messages are in the correct order and have correct content
for i, message in enumerate(data):
assert message["content"] == f"Test message {i}"
@ -195,8 +197,9 @@ async def test_create_batch_messages_limit(client, db_session, sample_data):
{
"content": f"Test message {i}",
"is_user": i % 2 == 0,
"metadata": {"batch_index": i}
} for i in range(101) # 101 messages
"metadata": {"batch_index": i},
}
for i in range(101) # 101 messages
]
}
@ -204,7 +207,7 @@ async def test_create_batch_messages_limit(client, db_session, sample_data):
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/batch",
json=test_messages,
)
assert response.status_code == 422 # Validation error
data = response.json()
assert "messages" in data["detail"][0]["loc"] # Error should mention messages field

98
tests/routes/test_metamessages.py
View File

@ -17,8 +17,9 @@ async def test_create_metamessage(client, db_session, sample_data):
await db_session.commit()
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/metamessages",
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
json={
"session_id": str(test_session.public_id),
"message_id": str(test_message.public_id),
"content": "Test Metamessage",
"metadata": {},
@ -27,6 +28,8 @@ async def test_create_metamessage(client, db_session, sample_data):
)
assert response.status_code == 200
data = response.json()
assert data["user_id"] == str(test_user.public_id)
assert data["session_id"] == str(test_session.public_id)
assert data["message_id"] == str(test_message.public_id)
assert data["content"] == "Test Metamessage"
assert data["metadata"] == {}
@ -46,19 +49,23 @@ async def test_get_metamessage(client, db_session, sample_data):
db_session.add(test_message)
await db_session.commit()
test_metamessage = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage",
metadata={},
h_metadata={},
metamessage_type="test_type",
)
db_session.add(test_metamessage)
await db_session.commit()
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/metamessages/{test_metamessage.public_id}/?message_id={test_message.public_id}"
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages/{test_metamessage.public_id}"
)
assert response.status_code == 200
data = response.json()
assert data["user_id"] == str(test_user.public_id)
assert data["session_id"] == str(test_session.public_id)
assert data["message_id"] == str(test_message.public_id)
assert data["content"] == "Test Metamessage"
assert data["metadata"] == {}
@ -66,7 +73,7 @@ async def test_get_metamessage(client, db_session, sample_data):
@pytest.mark.asyncio
async def test_get_metamessages(client, db_session, sample_data):
async def test_get_metamessages_by_session(client, db_session, sample_data):
test_app, test_user = sample_data
# Create a test session
test_session = models.Session(user_id=test_user.public_id)
@ -77,25 +84,35 @@ async def test_get_metamessages(client, db_session, sample_data):
)
db_session.add(test_message)
await db_session.commit()
# Create metamessages for the same session
test_metamessage_1 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_2 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_3 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_4 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage",
h_metadata={},
@ -107,25 +124,29 @@ async def test_get_metamessages(client, db_session, sample_data):
db_session.add(test_metamessage_4)
await db_session.commit()
# Filter by session and type
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/metamessages/list",
json={"metamessage_type": "test_type"},
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages/list",
json={
"session_id": str(test_session.public_id),
"metamessage_type": "test_type",
},
)
assert response.status_code == 200
data = response.json()
assert "items" in data
assert len(data["items"]) > 0
assert len(data["items"]) == 3
assert data["items"][0]["content"] == "Test Metamessage"
assert data["items"][0]["metamessage_type"] == "test_type"
assert data["items"][0]["session_id"] == str(test_session.public_id)
assert data["items"][0]["metadata"] == {}
@pytest.mark.asyncio
async def test_get_metamessage_by_user(client, db_session, sample_data):
test_app, test_user = sample_data
# Create a 3 test sessions
# Create 3 test sessions
test_session_1 = models.Session(user_id=test_user.public_id)
test_session_2 = models.Session(user_id=test_user.public_id)
test_session_3 = models.Session(user_id=test_user.public_id)
@ -149,48 +170,87 @@ async def test_get_metamessage_by_user(client, db_session, sample_data):
db_session.add(test_message_3)
await db_session.commit()
# create a metamessage on each message
# Create metamessages across different sessions
test_metamessage_1 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session_1.public_id,
message_id=test_message_1.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_2 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session_2.public_id,
message_id=test_message_2.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_3 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session_3.public_id,
message_id=test_message_3.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_4 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session_3.public_id,
message_id=test_message_3.public_id,
content="Test Metamessage",
h_metadata={},
metamessage_type="test_type_2",
)
# Create a user-level metamessage (no session/message)
test_metamessage_5 = models.Metamessage(
user_id=test_user.public_id,
content="User level metamessage",
h_metadata={},
metamessage_type="test_type",
)
db_session.add(test_metamessage_1)
db_session.add(test_metamessage_2)
db_session.add(test_metamessage_3)
db_session.add(test_metamessage_4)
db_session.add(test_metamessage_5)
await db_session.commit()
# Filter only by type across all user's metamessages
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages/list",
json={"metamessage_type": "test_type"},
)
assert response.status_code == 200
data = response.json()
assert len(data["items"]) > 0
assert len(data["items"]) == 3
assert data["items"][0]["content"] == "Test Metamessage"
assert len(data["items"]) == 4 # All test_type metamessages for the user
assert data["items"][0]["content"] in ["Test Metamessage", "User level metamessage"]
assert data["items"][0]["metamessage_type"] == "test_type"
assert data["items"][0]["metadata"] == {}
assert data["items"][0]["user_id"] == str(test_user.public_id)
@pytest.mark.asyncio
async def test_create_user_level_metamessage(client, db_session, sample_data):
test_app, test_user = sample_data
# Create a user-level metamessage (no session or message)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
json={
"content": "User level insight",
"metadata": {"source": "user_profile"},
"metamessage_type": "user_insight",
},
)
assert response.status_code == 200
data = response.json()
assert data["user_id"] == str(test_user.public_id)
assert data["session_id"] is None
assert data["message_id"] is None
assert data["content"] == "User level insight"
assert data["metadata"] == {"source": "user_profile"}
assert data["metamessage_type"] == "user_insight"
@pytest.mark.asyncio
@ -206,21 +266,27 @@ async def test_update_metamessage(client, db_session, sample_data):
db_session.add(test_message)
await db_session.commit()
test_metamessage = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage",
metadata={},
h_metadata={},
metamessage_type="test_type",
)
db_session.add(test_metamessage)
await db_session.commit()
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/metamessages/{test_metamessage.public_id}",
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages/{test_metamessage.public_id}",
json={
"message_id": str(test_message.public_id),
"metadata": {"new_key": "new_value"},
"metamessage_type": "updated_type",
},
)
assert response.status_code == 200
data = response.json()
assert data["metadata"] == {"new_key": "new_value"}
assert data["metamessage_type"] == "updated_type"
assert data["user_id"] == str(test_user.public_id)
assert data["session_id"] == str(test_session.public_id)
assert data["message_id"] == str(test_message.public_id)

628
tests/routes/test_scoped_api.py Normal file
View File

@ -0,0 +1,628 @@
from nanoid import generate as generate_nanoid
from src.security import JWTParams, create_jwt
def test_create_app_with_auth(auth_client):
name = str(generate_nanoid())
response = auth_client.post(
"/v1/apps", json={"name": name, "metadata": {"key": "value"}}
)
# Check expected behavior based on auth type
if auth_client.auth_type != "admin":
assert response.status_code == 401
return
assert response.status_code == 200
def test_auth_response_time(auth_client):
name = str(generate_nanoid())
import time
start_time = time.time()
response = auth_client.post(
"/v1/apps", json={"name": name, "metadata": {"key": "value"}}
)
end_time = time.time()
response_time = end_time - start_time
print(
f"Server response time for client {auth_client.auth_type}: {response_time:.6f} seconds"
)
# Check expected behavior based on auth type
if auth_client.auth_type != "admin":
assert response.status_code == 401
return
assert response.status_code == 200
def test_get_or_create_app_with_auth(auth_client):
name = str(generate_nanoid())
# Should return a ResourceNotFoundException with 404 status
response = auth_client.get(f"/v1/apps/name/{name}")
if auth_client.auth_type != "admin":
assert response.status_code == 401
return
assert response.status_code == 404
response = auth_client.get(f"/v1/apps/get_or_create/{name}")
if auth_client.auth_type != "admin":
assert response.status_code == 401
return
assert response.status_code == 200
def test_get_app_by_id_with_auth(auth_client, sample_data):
test_app, _ = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
response = auth_client.get(f"/v1/apps?app_id={test_app.public_id}")
# Admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
def test_get_app_from_token(auth_client, sample_data):
test_app, _ = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
response = auth_client.get("/v1/apps")
if auth_client.auth_type == "empty":
assert response.status_code == 200
assert response.json()["id"] == test_app.public_id
else:
assert response.status_code == 401
def test_get_app_by_name_with_auth(auth_client, sample_data):
test_app, _ = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
# Note that this will still fail because name route requires admin
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
response = auth_client.get(f"/v1/apps/name/{test_app.name}")
# Only admin JWT should be allowed
if auth_client.auth_type == "admin":
assert response.status_code == 200
else:
assert response.status_code == 401
def test_update_app_with_auth(auth_client, sample_data):
test_app, _ = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
# Note that this will still fail because name route requires admin
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
new_name = str(generate_nanoid())
response = auth_client.put(
f"/v1/apps/{test_app.public_id}",
json={"name": new_name, "metadata": {"new_key": "new_value"}},
)
# Only admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
def test_update_app_with_wrong_auth(auth_client, sample_data):
test_app, _ = sample_data
different_app = str(generate_nanoid())
if auth_client.auth_type == "empty":
# For non-admin, include the *wrong* app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=different_app))}"
)
new_name = str(generate_nanoid())
response = auth_client.put(
f"/v1/apps/{test_app.public_id}",
json={"name": new_name, "metadata": {"new_key": "new_value"}},
)
# Only admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type == "admin":
assert response.status_code == 200
else:
# wrong app_id should be rejected
assert response.status_code == 401
def test_create_user_with_auth(auth_client, sample_data):
test_app, _ = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
name = str(generate_nanoid())
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users",
json={"name": name, "metadata": {"user_key": "user_value"}},
)
# Only admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
def test_get_user_by_id_with_auth(auth_client, sample_data):
test_app, test_user = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}"
)
# Admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
# Test with user-scoped JWT
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
)
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}"
)
assert response.status_code == 200
response2 = auth_client.get(f"/v1/apps/{test_app.public_id}/users")
assert response2.status_code == 200
print(response2.json())
assert response2.json()["id"] == test_user.public_id
def test_get_user_by_name_with_auth(auth_client, sample_data):
test_app, test_user = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/name/{test_user.name}"
)
# Admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
def test_update_user_with_auth(auth_client, sample_data):
test_app, test_user = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
new_name = str(generate_nanoid())
response = auth_client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}",
json={"name": new_name, "metadata": {"updated_key": "updated_value"}},
)
# Admin JWT or JWT with matching app_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
# Test with user-scoped JWT
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
)
response = auth_client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}",
json={
"name": str(generate_nanoid()),
"metadata": {"user_key": "user_value"},
},
)
assert response.status_code == 200
def test_create_session_with_auth(auth_client, sample_data):
test_app, test_user = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id and user_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={},
)
# Only admin JWT or JWT with matching app_id and user_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
# Remove app_id from header and make sure user-scoped key works too
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
)
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={},
)
assert response.status_code == 200
def test_get_session_by_id_with_auth(auth_client, sample_data):
test_app, test_user = sample_data
# First create a session
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
create_response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={},
)
if auth_client.auth_type not in ["admin", "empty"]:
assert create_response.status_code == 401
return
assert create_response.status_code == 200
session_id = create_response.json()["id"]
# Test with app and user scoped JWT
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={session_id}"
)
assert response.status_code == 200
if auth_client.auth_type == "empty":
# Test with session-scoped JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(se=session_id))}"
)
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={session_id}"
)
assert response.status_code == 200
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions"
)
assert response.status_code == 200
assert response.json()["id"] == session_id
# Test with wrong session_id (should be 401 since we have a session-scoped JWT)
assert (
auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={generate_nanoid()}"
).status_code
== 401
)
# Test with user-scoped JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
)
assert (
auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={session_id}"
).status_code
== 200
)
assert (
auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions"
).status_code
== 401
)
# Test with app-scoped JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
assert (
auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={session_id}"
).status_code
== 200
)
assert (
auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions"
).status_code
== 401
)
# Test with wrong session_id (should be 404 since we have an app-scoped JWT)
assert (
auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={generate_nanoid()}"
).status_code
== 404
)
def test_create_collection(auth_client, sample_data) -> None:
test_app, test_user = sample_data
if auth_client.auth_type == "empty":
# For non-admin, include the app_id and user_id in the JWT
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection", "metadata": {}},
)
# Only admin JWT or JWT with matching app_id and user_id should be allowed
if auth_client.auth_type in ["admin", "empty"]:
assert response.status_code == 200
else:
assert response.status_code == 401
# Remove app_id from header and make sure user-scoped key works too
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
)
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection2", "metadata": {}},
)
assert response.status_code == 200
# Remove user_id from header and make sure app-scoped key works too
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
)
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection3", "metadata": {}},
)
assert response.status_code == 200
def test_get_collection_by_id_with_auth(auth_client, sample_data) -> None:
test_app, test_user = sample_data
# First create a collection
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
create_response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection_get", "metadata": {}},
)
if auth_client.auth_type not in ["admin", "empty"]:
assert create_response.status_code == 401
return
assert create_response.status_code == 200
collection_id = create_response.json()["id"]
# Test with app and user scoped JWT
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections?collection_id={collection_id}"
)
assert response.status_code == 200
# Test with collection-scoped JWT
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(co=collection_id))}"
)
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections?collection_id={collection_id}"
)
assert response.status_code == 200
# Test auto resolution of ID
response2 = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections"
)
assert response2.status_code == 200
assert response2.json()["id"] == collection_id
def test_get_collection_by_name_with_auth(auth_client, sample_data) -> None:
test_app, test_user = sample_data
collection_name = f"test_collection_{generate_nanoid()}"
# First create a collection
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
create_response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": collection_name, "metadata": {}},
)
if auth_client.auth_type not in ["admin", "empty"]:
assert create_response.status_code == 401
return
assert create_response.status_code == 200
# Test with app and user scoped JWT
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/name/{collection_name}"
)
assert response.status_code == 200
def test_create_document_with_auth(auth_client, sample_data) -> None:
test_app, test_user = sample_data
# First create a collection
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
create_collection_response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection_docs", "metadata": {}},
)
if auth_client.auth_type not in ["admin", "empty"]:
assert create_collection_response.status_code == 401
return
assert create_collection_response.status_code == 200
collection_id = create_collection_response.json()["id"]
# Create document with app and user scoped JWT
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents",
json={"content": "Test document content", "metadata": {"doc_key": "doc_value"}},
)
assert response.status_code == 200
# Test with collection-scoped JWT
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(co=collection_id))}"
)
response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents",
json={"content": "Test document with collection JWT", "metadata": {}},
)
assert response.status_code == 200
def test_get_document_with_auth(auth_client, sample_data) -> None:
test_app, test_user = sample_data
# First create a collection and document
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
)
create_collection_response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "test_collection_get_doc", "metadata": {}},
)
if auth_client.auth_type not in ["admin", "empty"]:
assert create_collection_response.status_code == 401
return
assert create_collection_response.status_code == 200
collection_id = create_collection_response.json()["id"]
create_doc_response = auth_client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents",
json={"content": "Test document for retrieval", "metadata": {}},
)
assert create_doc_response.status_code == 200
document_id = create_doc_response.json()["id"]
# Get document with app and user scoped JWT
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/{document_id}"
)
assert response.status_code == 200
# Test with collection-scoped JWT
if auth_client.auth_type == "empty":
auth_client.headers["Authorization"] = (
f"Bearer {create_jwt(JWTParams(co=collection_id))}"
)
response = auth_client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/{document_id}"
)
assert response.status_code == 200

26
tests/routes/test_sessions.py
View File

@ -118,7 +118,7 @@ async def test_delete_session(client, db_session, sample_data):
)
assert response.status_code == 200
response = client.get(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}"
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={test_session.public_id}"
)
data = response.json()
assert data["is_active"] is False
@ -266,24 +266,32 @@ async def test_deep_clone_session(client, db_session, sample_data):
await db_session.commit()
test_metamessage_1 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage 1",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_2 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage 2",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_3 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message2.public_id,
content="Test Metamessage 3",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_4 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message2.public_id,
content="Test Metamessage 4",
h_metadata={},
@ -325,8 +333,8 @@ async def test_deep_clone_session(client, db_session, sample_data):
assert data["items"][1]["metadata"] == {"key": "value2"}
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{cloned_session_id}/metamessages/list",
json={},
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages/list",
json={"session_id": cloned_session_id},
)
assert response.status_code == 200
@ -374,24 +382,32 @@ async def test_partial_deep_clone_session(client, db_session, sample_data):
await db_session.commit()
test_metamessage_1 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage 1",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_2 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message.public_id,
content="Test Metamessage 2",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_3 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message2.public_id,
content="Test Metamessage 3",
h_metadata={},
metamessage_type="test_type",
)
test_metamessage_4 = models.Metamessage(
user_id=test_user.public_id,
session_id=test_session.public_id,
message_id=test_message2.public_id,
content="Test Metamessage 4",
h_metadata={},
@ -429,8 +445,8 @@ async def test_partial_deep_clone_session(client, db_session, sample_data):
assert data["items"][0]["metadata"] == {"key": "value"}
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{cloned_session_id}/metamessages/list",
json={},
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages/list",
json={"session_id": cloned_session_id},
)
assert response.status_code == 200

4
tests/routes/test_users.py
View File

@ -17,7 +17,9 @@ def test_create_user(client, sample_data):
def test_get_user_by_id(client, sample_data):
test_app, test_user = sample_data
response = client.get(f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}")
response = client.get(
f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}"
)
assert response.status_code == 200
data = response.json()
assert data["name"] == test_user.name

190
tests/routes/test_validation_api.py
View File

@ -1,4 +1,3 @@
import pytest
from nanoid import generate as generate_nanoid
@ -32,8 +31,7 @@ def test_user_validations_api(client, sample_data):
# Test name too short
response = client.post(
f"/v1/apps/{test_app.public_id}/users",
json={"name": "", "metadata": {}}
f"/v1/apps/{test_app.public_id}/users", json={"name": "", "metadata": {}}
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -43,8 +41,7 @@ def test_user_validations_api(client, sample_data):
# Test name too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users",
json={"name": "a" * 101, "metadata": {}}
f"/v1/apps/{test_app.public_id}/users", json={"name": "a" * 101, "metadata": {}}
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -58,18 +55,14 @@ def test_message_validations_api(client, sample_data):
# Create a test session first
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
# Test content too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages",
json={
"content": "a" * 50001,
"is_user": True,
"metadata": {}
}
json={"content": "a" * 50001, "is_user": True, "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -80,11 +73,7 @@ def test_message_validations_api(client, sample_data):
# Test invalid is_user type
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages",
json={
"content": "test",
"is_user": "not a bool",
"metadata": {}
}
json={"content": "test", "is_user": "not a bool", "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -98,7 +87,7 @@ def test_collection_validations_api(client, sample_data):
# Test name too short
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "", "metadata": {}}
json={"name": "", "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -109,7 +98,7 @@ def test_collection_validations_api(client, sample_data):
# Test name too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "a" * 101, "metadata": {}}
json={"name": "a" * 101, "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -120,7 +109,7 @@ def test_collection_validations_api(client, sample_data):
# Test 'honcho' name restriction
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": "honcho", "metadata": {}}
json={"name": "honcho", "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -134,14 +123,14 @@ def test_document_validations_api(client, sample_data):
# Create a collection first
collection_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": str(generate_nanoid()), "metadata": {}}
json={"name": str(generate_nanoid()), "metadata": {}},
)
collection_id = collection_response.json()["id"]
# Test content too short
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents",
json={"content": "", "metadata": {}}
json={"content": "", "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -152,7 +141,7 @@ def test_document_validations_api(client, sample_data):
# Test content too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents",
json={"content": "a" * 100001, "metadata": {}}
json={"content": "a" * 100001, "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -166,14 +155,14 @@ def test_document_query_validations_api(client, sample_data):
# Create a collection first
collection_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": str(generate_nanoid()), "metadata": {}}
json={"name": str(generate_nanoid()), "metadata": {}},
)
collection_id = collection_response.json()["id"]
# Test query too short
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/query",
json={"query": "", "top_k": 5}
json={"query": "", "top_k": 5},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -184,7 +173,7 @@ def test_document_query_validations_api(client, sample_data):
# Test query too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/query",
json={"query": "a" * 1001, "top_k": 5}
json={"query": "a" * 1001, "top_k": 5},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -195,7 +184,7 @@ def test_document_query_validations_api(client, sample_data):
# Test top_k too small
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/query",
json={"query": "test", "top_k": 0}
json={"query": "test", "top_k": 0},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -206,7 +195,7 @@ def test_document_query_validations_api(client, sample_data):
# Test top_k too large
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/query",
json={"query": "test", "top_k": 51}
json={"query": "test", "top_k": 51},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -220,23 +209,19 @@ def test_message_batch_validations_api(client, sample_data):
# Create a test session first
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
# Test batch too large
messages = [
{
"content": f"test message {i}",
"is_user": True,
"metadata": {}
}
{"content": f"test message {i}", "is_user": True, "metadata": {}}
for i in range(101) # Create 101 messages
]
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages/batch",
json={"messages": messages}
json={"messages": messages},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -250,25 +235,26 @@ def test_metamessage_validations_api(client, sample_data):
# Create session and message first
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
message_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages",
json={"content": "test message", "is_user": True, "metadata": {}}
json={"content": "test message", "is_user": True, "metadata": {}},
)
message_id = message_response.json()["id"]
# Test metamessage_type too short
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/metamessages",
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
json={
"metamessage_type": "",
"content": "test content",
"session_id": session_id,
"message_id": message_id,
"metadata": {}
}
"metadata": {},
},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -278,13 +264,14 @@ def test_metamessage_validations_api(client, sample_data):
# Test metamessage_type too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/metamessages",
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
json={
"metamessage_type": "a" * 51,
"content": "test content",
"session_id": session_id,
"message_id": message_id,
"metadata": {}
}
"metadata": {},
},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -294,13 +281,14 @@ def test_metamessage_validations_api(client, sample_data):
# Test content too long
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/metamessages",
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
json={
"metamessage_type": "test_type",
"content": "a" * 50001,
"message_id": message_id,
"metadata": {}
}
"session_id": session_id,
"metadata": {},
},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -314,14 +302,14 @@ def test_collection_update_validations_api(client, sample_data):
# Create a collection first
collection_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": str(generate_nanoid()), "metadata": {}}
json={"name": str(generate_nanoid()), "metadata": {}},
)
collection_id = collection_response.json()["id"]
# Test honcho name in update
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}",
json={"name": "honcho", "metadata": {}}
json={"name": "honcho", "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -335,26 +323,26 @@ def test_document_update_validations_api(client, sample_data):
# Create collection and document first
collection_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"name": str(generate_nanoid()), "metadata": {}}
json={"name": str(generate_nanoid()), "metadata": {}},
)
collection_id = collection_response.json()["id"]
document_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents",
json={"content": "test content", "metadata": {}}
json={"content": "test content", "metadata": {}},
)
document_id = document_response.json()["id"]
# Test content too long in update
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection_id}/documents/{document_id}",
json={"content": "a" * 100001, "metadata": {}}
json={"content": "a" * 100001, "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
assert error["loc"] == ["body", "content"]
assert error["msg"] == "String should have at most 100000 characters"
assert error["type"] == "string_too_long"
assert error["type"] == "string_too_long"
def test_session_validations_api(client, sample_data):
@ -362,14 +350,14 @@ def test_session_validations_api(client, sample_data):
# Create a test session first
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
# Test invalid metadata type
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}",
json={"metadata": "not a dict"}
json={"metadata": "not a dict"},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -379,31 +367,89 @@ def test_session_validations_api(client, sample_data):
# Test empty update
response = client.put(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}",
json={}
json={},
)
assert response.status_code == 422
def test_agent_query_validations_api(client, sample_data):
def test_agent_query_validations_api(client, sample_data, monkeypatch):
# Mock the functions in agent.py that are causing the database issues
# Create a mock collection with a public_id
class MockCollection:
def __init__(self):
self.public_id = "mock_collection_id"
# Mock collection retrieval/creation function
async def mock_get_or_create_collection(*args, **kwargs):
return MockCollection()
async def mock_chat_history(*args, **kwargs):
return "Mock chat history", [], []
async def mock_get_long_term_facts(*args, **kwargs):
return ["Mock fact 1", "Mock fact 2"]
async def mock_run_tom_inference(*args, **kwargs):
return "Mock TOM inference"
async def mock_generate_user_representation(*args, **kwargs):
return "Mock user representation"
# Mock the Dialectic.call method
async def mock_dialectic_call(self):
# Create a mock response that will work with line 300 in agent.py:
# return schemas.AgentChat(content=response[0]["text"])
return [{"text": "Mock response"}]
# Mock the Dialectic.stream method
def mock_dialectic_stream(self):
class MockStream:
def __enter__(self):
return self
def __exit__(self, *args):
pass
@property
def text_stream(self):
yield "Mock streamed response"
return MockStream()
# Apply the monkeypatches
monkeypatch.setattr(
"src.crud.get_or_create_user_protected_collection",
mock_get_or_create_collection,
)
monkeypatch.setattr("src.utils.history.get_summarized_history", mock_chat_history)
monkeypatch.setattr("src.agent.get_long_term_facts", mock_get_long_term_facts)
monkeypatch.setattr("src.agent.run_tom_inference", mock_run_tom_inference)
monkeypatch.setattr(
"src.agent.generate_user_representation", mock_generate_user_representation
)
monkeypatch.setattr("src.agent.Dialectic.call", mock_dialectic_call)
monkeypatch.setattr("src.agent.Dialectic.stream", mock_dialectic_stream)
test_app, test_user = sample_data
# Create a session first since agent queries are likely session-based
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
# Test valid string query (under 10000 chars)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/chat",
json={"queries": "a" * 9999}
json={"queries": "a" * 9999},
)
assert response.status_code == 200
# Test string query too long (over 10000 chars)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/chat",
json={"queries": "a" * 10001}
json={"queries": "a" * 10001},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -414,14 +460,14 @@ def test_agent_query_validations_api(client, sample_data):
# Test valid list query (under 25 items, each under 10000 chars)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/chat",
json={"queries": ["a" * 9999 for _ in range(25)]}
json={"queries": ["a" * 9999 for _ in range(25)]},
)
assert response.status_code == 200
# Test list too long (over 25 items)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/chat",
json={"queries": ["test" for _ in range(26)]}
json={"queries": ["test" for _ in range(26)]},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -431,7 +477,7 @@ def test_agent_query_validations_api(client, sample_data):
# Test list item too long (item over 10000 chars)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/chat",
json={"queries": ["a" * 10001]}
json={"queries": ["a" * 10001]},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -442,7 +488,7 @@ def test_agent_query_validations_api(client, sample_data):
# Test that strings over 20 chars are allowed
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/chat",
json={"queries": "a" * 100} # 100 chars should be fine
json={"queries": "a" * 100}, # 100 chars should be fine
)
assert response.status_code == 200
@ -451,14 +497,14 @@ def test_required_field_validations_api(client, sample_data):
test_app, test_user = sample_data
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
# Test missing required content in message
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages",
json={"is_user": True, "metadata": {}}
json={"is_user": True, "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -468,7 +514,7 @@ def test_required_field_validations_api(client, sample_data):
# Test missing required is_user in message
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages",
json={"content": "test", "metadata": {}}
json={"content": "test", "metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -478,7 +524,7 @@ def test_required_field_validations_api(client, sample_data):
# Test missing required name in collection
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
json={"metadata": {}}
json={"metadata": {}},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -491,14 +537,14 @@ def test_filter_validations_api(client, sample_data):
# Create a session first
session_response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
json={"metadata": {}}
json={"metadata": {}},
)
session_id = session_response.json()["id"]
# Test invalid filter type in message list (at session level)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}/messages/list",
json={"filter": "not a dict"}
json={"filter": "not a dict"},
)
assert response.status_code == 422
error = response.json()["detail"][0]
@ -508,7 +554,7 @@ def test_filter_validations_api(client, sample_data):
# Test invalid filter type in collection list (at user level)
response = client.post(
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/list",
json={"filter": "not a dict"}
json={"filter": "not a dict"},
)
assert response.status_code == 422
error = response.json()["detail"][0]

13
tests/test_schema_validations.py
View File

@ -1,14 +1,15 @@
import pytest
from pydantic import ValidationError
from src.schemas import (
AppCreate,
UserCreate,
MessageCreate,
MetamessageCreate,
CollectionCreate,
DocumentCreate,
DocumentQuery,
MessageBatchCreate,
MessageCreate,
MetamessageCreate,
UserCreate,
)
@ -200,9 +201,7 @@ class TestDocumentQueryValidations:
class TestMessageBatchValidations:
def test_valid_message_batch(self):
batch = MessageBatchCreate(
messages=[
MessageCreate(content="test", is_user=True, metadata={})
]
messages=[MessageCreate(content="test", is_user=True, metadata={})]
)
assert len(batch.messages) == 1
@ -215,4 +214,4 @@ class TestMessageBatchValidations:
]
)
error_dict = exc_info.value.errors()[0]
assert error_dict["type"] == "too_long"
assert error_dict["type"] == "too_long"

191
tests/utils/test_model_client.py Normal file
View File

@ -0,0 +1,191 @@
import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.utils.model_client import (
DEFAULT_MAX_TOKENS,
DEFAULT_MODELS,
DEFAULT_TEMPERATURE,
ModelClient,
ModelProvider,
)
@pytest.fixture
def mock_env(monkeypatch):
monkeypatch.setenv("ANTHROPIC_API_KEY", "mock-anthropic-api-key")
monkeypatch.setenv("OPENAI_COMPATIBLE_API_KEY", "mock-openai-api-key")
monkeypatch.setenv("OPENAI_API_KEY", "mock-openai-api-key")
monkeypatch.setenv("GROQ_API_KEY", "mock-groq-api-key")
monkeypatch.setenv("CEREBRAS_API_KEY", "mock-cerebras-api-key")
monkeypatch.setenv("OPENROUTER_API_KEY", "mock-openrouter-api-key")
# Test fixtures
@pytest.fixture
def mock_anthropic_response():
mock_response = MagicMock()
mock_content = MagicMock()
mock_content.type = "text"
mock_content.text = "Test response"
mock_response.content = [mock_content]
return mock_response
@pytest.fixture
def mock_openai_response():
mock_response = MagicMock()
mock_message = MagicMock()
mock_message.content = "Test response"
mock_choice = MagicMock()
mock_choice.message = mock_message
mock_response.choices = [mock_choice]
return mock_response
@pytest.fixture
def mock_anthropic_client():
with patch("src.utils.model_client.AsyncAnthropic") as mock:
client = MagicMock()
client.messages.create = AsyncMock()
client.messages.stream = AsyncMock()
mock.return_value = client
yield client
@pytest.fixture
def mock_openai_client():
with patch("src.utils.model_client.AsyncOpenAI") as mock:
client = MagicMock()
client.chat.completions.create = AsyncMock()
mock.return_value = client
yield client
# Initialization Tests
def test_default_initialization(mock_env):
"""Test default initialization with Anthropic provider."""
client = ModelClient()
assert client.provider == ModelProvider.ANTHROPIC
assert client.model == DEFAULT_MODELS[ModelProvider.ANTHROPIC]
assert client.base_url is None
def test_custom_model_initialization():
"""Test initialization with custom model name."""
custom_model = "custom-model"
client = ModelClient(model=custom_model)
assert client.model == custom_model
def test_custom_api_key_initialization():
"""Test initialization with custom API key."""
custom_key = "test-api-key"
client = ModelClient(api_key=custom_key)
assert client.api_key == custom_key
def test_custom_base_url_initialization():
"""Test initialization with custom base URL."""
custom_url = "https://custom-api.example.com"
client = ModelClient(base_url=custom_url)
assert client.base_url == custom_url
def test_unsupported_provider_initialization():
"""Test initialization with unsupported provider."""
with pytest.raises(ValueError, match="is not a valid ModelProvider"):
ModelClient(provider=ModelProvider("unsupported"))
def test_missing_api_key_initialization():
"""Test initialization without required API key."""
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="API key is required"):
ModelClient()
# Message Creation Tests
def test_create_message():
"""Test message creation with different roles."""
client = ModelClient()
message = client.create_message("user", "Hello")
assert message == {"role": "user", "content": "Hello"}
# Generation Tests
@pytest.mark.asyncio
async def test_generate_anthropic(mock_anthropic_client, mock_anthropic_response):
"""Test generation with Anthropic provider."""
mock_anthropic_client.messages.create.return_value = mock_anthropic_response
client = ModelClient(provider=ModelProvider.ANTHROPIC)
messages = [{"role": "user", "content": "Hello"}]
response = await client.generate(messages)
assert response == "Test response"
# Verify the API call
mock_anthropic_client.messages.create.assert_called_once()
call_args = mock_anthropic_client.messages.create.call_args[1]
assert call_args["model"] == DEFAULT_MODELS[ModelProvider.ANTHROPIC]
assert call_args["messages"] == messages
assert call_args["max_tokens"] == DEFAULT_MAX_TOKENS
assert call_args["temperature"] == DEFAULT_TEMPERATURE
@pytest.mark.asyncio
async def test_generate_openai(mock_openai_client, mock_openai_response, mock_env):
"""Test generation with OpenAI provider."""
mock_openai_client.chat.completions.create.return_value = mock_openai_response
client = ModelClient(provider=ModelProvider.OPENAI)
messages = [{"role": "user", "content": "Hello"}]
response = await client.generate(messages)
assert response == "Test response"
# Verify the API call
mock_openai_client.chat.completions.create.assert_called_once()
call_args = mock_openai_client.chat.completions.create.call_args[1]
assert call_args["model"] == DEFAULT_MODELS[ModelProvider.OPENAI]
assert call_args["messages"] == messages
assert call_args["max_tokens"] == DEFAULT_MAX_TOKENS
assert call_args["temperature"] == DEFAULT_TEMPERATURE
@pytest.mark.asyncio
async def test_generate_with_system_prompt(
mock_anthropic_client, mock_anthropic_response
):
"""Test generation with system prompt."""
mock_anthropic_client.messages.create.return_value = mock_anthropic_response
client = ModelClient(provider=ModelProvider.ANTHROPIC)
messages = [{"role": "user", "content": "Hello"}]
system = "You are a helpful assistant"
response = await client.generate(messages, system=system)
assert response == "Test response"
# Verify the API call
mock_anthropic_client.messages.create.assert_called_once()
call_args = mock_anthropic_client.messages.create.call_args[1]
assert call_args["system"] == system
@pytest.mark.asyncio
async def test_generate_with_caching(mock_anthropic_client, mock_anthropic_response):
"""Test generation with caching enabled."""
mock_anthropic_client.messages.create.return_value = mock_anthropic_response
client = ModelClient(provider=ModelProvider.ANTHROPIC)
messages = [{"role": "user", "content": "Hello"}]
system = "You are a helpful assistant"
response = await client.generate(messages, system=system, use_caching=True)
assert response == "Test response"
# Verify the API call
mock_anthropic_client.messages.create.assert_called_once()
call_args = mock_anthropic_client.messages.create.call_args[1]
assert call_args["system"] == [
{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}
]

1416
uv.lock
View File

File diff suppressed because it is too large Load Diff