Migrate to Peer Paradigm (#131)
* Initial Model Changes * fix migration * update schemas * handle router changes * make name FK and corresponding crud changes * fix routers * comment metamessage references * add bulk peer session operations * update messages router * fix require_auth to make app runnable * remove peer from get messages * add new routes * implement new crud methods for session peers * alter keys router * add feature flags dict and token limit + fix SessionContext * fix: paginate get_session_peers and make tokens/summary query params in get_session_context * feat: add create_messages_for_peer, get_messages_for_peer * fix: make session_peers a Table * finalize upgrade * fix: working migration * fixes: schemas, crud, routes * add token count * fix migration errors discovered from db with data in it * fixes: unify with sdk * downgrade * feat: swap jwts to new paradigm * fix unit tests * fix tests pt 2 * fix: handle foreign key errors in create_messages * fix downgrade * downgrade queue changes * feat: add search to resources, make get_messages handle limits, add get_representation to peer * chore: beef up tests * fix: move chat and rep params to post body, add target to get_representation * fix get_user_protected_collection and embedding store * feat: add peer config to models, crud, schemas, routes * fix: update tests and fix list(tuple()) to dict() * add session peer left_at/joined_at and modify enqueue * [wip]: feat: refactor history to match new paradigm and implement get_context * fix messages enqueue and test it * chore: align deriver and new honcho paradigm * chore: update consumer * chore: get rid of is_user * feat: change queue tables to new key strat * fix: convert queue session_id to str properly * fix downgrade migration * feat: re-integrate old deriver * chore: coderabbit review, lots of small bug fixes * fix: fix batch migration of messages and token count * fix: mock ModelClient * CodeRabbit comments * CR comments 2 * fix: handle metadata and feature flags properly in get_or_creates * cr comments 3 * feature flag to configuration * feat: add real get crud * fix: remove reverse param from places it does not belong * add session.name constraint; narrow task type; disable deriver from configuration * get_or_add_peers_to_session + session peers limit * fix: add internal_metadata, fix agent * fix: move working rep into crud get/set, unstub get_working_representation * fix: don't payload metadata * peer protected collection -> global / local rep collections * fix: remove spurious mockery * feat: add english language search index * fix: remove spurious error * chore: 2.0.0 -- update readme, changelog, claude.md * chore: update readme for peer paradigm --------- Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Co-authored-by: Rajat Ahuja <rahuja445@gmail.com>
This commit is contained in:
parent
4fab671e20
commit
657feacc53
39
CHANGELOG.md
39
CHANGELOG.md
|
|
@ -5,6 +5,45 @@ 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/).
|
||||
|
||||
## [2.0.0]
|
||||
|
||||
### Added
|
||||
|
||||
- Ability to get a peer's working representation
|
||||
- Metadata to all data primitives (Workspaces, Peers, Sessions, Messages)
|
||||
- Internal metadata to store Honcho's state no longer exposed in API
|
||||
- Batch message operations and enhanced message querying with token and message count limits
|
||||
- Search and summary functionalities scoped by workspace, peer, and session
|
||||
- Session context retrieval with summaries and token allocation
|
||||
|
||||
### Changed
|
||||
|
||||
- New architecture centered around the concept of a "peer" replaces the former
|
||||
"app"/"user"/"session" paradigm
|
||||
- Workspaces replace "apps" as top-level namespace
|
||||
- Peers replace "users"
|
||||
- Sessions no longer nested beneath peers and no longer limited to a single
|
||||
user-assistant model. A session exists independently of any one peer and
|
||||
peers can be added to and removed from sessions.
|
||||
- Dialectic API is now part of the Peer, not the Session
|
||||
- Dialectic API now allows queries to be scoped to a session or "targeted"
|
||||
to a fellow peer
|
||||
- Database schema migrated to adopt workspace/peer/session naming and structure
|
||||
- Authentication and JWT scopes updated to workspace/peer/session hierarchy
|
||||
- Queue processing now works on 'work units' instead of sessions
|
||||
- Message token counting updated with tiktoken integration and fallback heuristic
|
||||
- Queue and message processing updated to handle sender/target and task types for multi-peer scenarios
|
||||
|
||||
### Fixed
|
||||
|
||||
- Improved error handling and validation for batch message operations and metadata
|
||||
|
||||
### Removed
|
||||
|
||||
- Metamessages removed in favor of metadata
|
||||
- Collections and Documents no longer exposed in the API, solely internal
|
||||
- Obsolete tests for apps, users, collections, documents, and metamessages
|
||||
|
||||
## [1.1.0]
|
||||
|
||||
### Added
|
||||
|
|
|
|||
100
CLAUDE.md
100
CLAUDE.md
|
|
@ -7,9 +7,64 @@ Honcho is an infrastructure layer for building AI agents with social cognition a
|
|||
- Personalizing user experiences through understanding user psychology
|
||||
- Providing a Dialectic API that injects personal context just-in-time
|
||||
- Supporting development of LLM-powered applications that adapt to end users
|
||||
- Enabling multi-peer sessions where multiple participants (users or agents) can interact
|
||||
|
||||
Honcho leverages the inherent theory-of-mind capabilities of LLMs to build coherent models of user psychology over time, enabling more personalized and effective AI interactions.
|
||||
|
||||
## Core Concepts
|
||||
|
||||
### Peer Paradigm
|
||||
Honcho uses a peer-based model where both users and agents are represented as "peers". This unified approach enables:
|
||||
- Multi-participant sessions with mixed human and AI agents
|
||||
- Configurable observation settings (which peers observe which others)
|
||||
- Flexible identity management for all participants
|
||||
|
||||
### Key Primitives
|
||||
- **Workspace** (formerly App): The root organizational unit containing all resources
|
||||
- **Peer** (formerly User): Any participant in the system (human or AI)
|
||||
- **Session**: A conversation context that can involve multiple peers
|
||||
- **Message**: Data units that can represent communication between peers OR arbitrary data ingested by a peer to enhance its global representation
|
||||
- **Collections & Documents**: Internal vector storage for theory-of-mind representations (not exposed via API)
|
||||
|
||||
## Architecture Overview
|
||||
|
||||
### API Structure
|
||||
All API routes follow the pattern: `/v1/{resource}/{id}/{action}`
|
||||
- **Workspaces**: Create, list, update, search
|
||||
- **Peers**: Create, list, update, chat (dialectic), messages, representation
|
||||
- **Sessions**: Create, list, update, delete, clone, manage peers, get context
|
||||
- **Messages**: Create (batch up to 100), list, get, update
|
||||
- **Keys**: Create scoped JWT tokens
|
||||
|
||||
### Key Features
|
||||
|
||||
#### Dialectic API (`/peers/{peer_id}/chat`)
|
||||
- Provides theory-of-mind informed responses
|
||||
- Integrates long-term facts from vector storage
|
||||
- Supports streaming responses
|
||||
- Configurable LLM providers
|
||||
|
||||
#### Message Processing Pipeline
|
||||
1. Messages created via API (batch or single)
|
||||
2. Enqueued for background processing:
|
||||
- `representation`: Update peer's theory of mind
|
||||
- `summary`: Create session summaries
|
||||
3. Session-based queue processing ensures order
|
||||
4. Results stored internally in vector DB
|
||||
|
||||
#### Theory of Mind System
|
||||
- Multiple implementation methods (conversational, single_prompt, long_term)
|
||||
- Facts extracted from messages and stored in collections
|
||||
- Representations combine short-term inference with long-term facts
|
||||
- Configurable via peer and session feature flags
|
||||
|
||||
### Configuration
|
||||
- Hierarchical config: config.toml + environment variables
|
||||
- Database settings with connection pooling
|
||||
- Multiple LLM provider support
|
||||
- Background worker (deriver) settings
|
||||
- Authentication can be toggled on/off
|
||||
|
||||
## Development Guide
|
||||
|
||||
### Commands
|
||||
|
|
@ -29,12 +84,51 @@ Honcho leverages the inherent theory-of-mind capabilities of LLMs to build coher
|
|||
- Docstrings: Use Google style docstrings
|
||||
|
||||
### Project Structure
|
||||
- FastAPI routes in src/routers/
|
||||
- SQLAlchemy ORM models in src/models.py with proper type annotations
|
||||
- Pydantic schemas in src/schemas.py for API validation
|
||||
```
|
||||
src/
|
||||
├── main.py # FastAPI app setup with middleware and exception handlers
|
||||
├── models.py # SQLAlchemy ORM models with proper type annotations
|
||||
├── schemas.py # Pydantic validation schemas for API
|
||||
├── crud.py # Database operations
|
||||
├── dependencies.py # Dependency injection (DB sessions)
|
||||
├── exceptions.py # Custom exception types
|
||||
├── security.py # JWT authentication
|
||||
├── agent.py # Dialectic API implementation
|
||||
├── routers/ # API endpoints
|
||||
│ ├── workspaces.py
|
||||
│ ├── peers.py
|
||||
│ ├── sessions.py
|
||||
│ ├── messages.py
|
||||
│ └── keys.py
|
||||
├── deriver/ # Background processing system
|
||||
│ ├── consumer.py # Message processing logic
|
||||
│ ├── queue.py # Queue management
|
||||
│ └── tom/ # Theory of Mind implementations
|
||||
└── utils/ # Utilities
|
||||
├── history.py # Session history management
|
||||
├── cache.py # Caching utilities
|
||||
└── model_client.py # LLM client abstraction
|
||||
```
|
||||
- Tests in pytest with fixtures in tests/conftest.py
|
||||
- Use environment variables via python-dotenv (.env)
|
||||
|
||||
### Database Design
|
||||
- All tables use text IDs (nanoid format) as primary keys
|
||||
- Composite foreign keys for multi-tenant relationships
|
||||
- Feature flags on workspace, peer, and session levels
|
||||
- Token counting on messages for usage tracking
|
||||
- JSONB metadata fields for extensibility
|
||||
- HNSW indexes for vector similarity search
|
||||
|
||||
### Key Architectural Decisions
|
||||
1. **Multi-Peer Sessions**: Sessions can have multiple participants with different observation settings
|
||||
2. **Flexible Theory of Mind**: Pluggable ToM implementations (conversational, single_prompt, long_term)
|
||||
3. **Background Processing**: Async queue system for expensive operations
|
||||
4. **Provider Abstraction**: Model client supports multiple LLM providers
|
||||
5. **Scoped Authentication**: JWT tokens can be scoped to workspace, peer, or session level
|
||||
6. **Batch Operations**: Support for bulk message creation (up to 100 messages)
|
||||
7. **Session History**: Two-tier summarization (short every 20 messages, long every 60)
|
||||
|
||||
### Error Handling
|
||||
- Custom exceptions defined in src/exceptions.py
|
||||
- Use specific exception types (ResourceNotFoundException, ValidationException, etc.)
|
||||
|
|
|
|||
136
README.md
136
README.md
|
|
@ -1,6 +1,6 @@
|
|||
# 🫡 Honcho
|
||||
|
||||

|
||||

|
||||
[](https://discord.gg/plasticlabs)
|
||||
[](https://arxiv.org/abs/2310.06983)
|
||||

|
||||
|
|
@ -9,9 +9,7 @@
|
|||
[](https://pypi.org/project/honcho-ai/)
|
||||
[](https://npmjs.org/package/honcho-ai)
|
||||
|
||||
Honcho is a platform for making AI agents and LLM powered applications that are personalized
|
||||
to their end users. It leverages the inherent theory-of-mind capabilities of
|
||||
LLMs to cohere to user psychology over time.
|
||||
Honcho is an infrastructure layer for building AI agents with social cognition and theory-of-mind capabilities. It enables developers to create AI agents and LLM-powered applications that are personalized to their end users by leveraging the inherent theory-of-mind capabilities of LLMs to build coherent models of user psychology over time.
|
||||
|
||||
Read about the project [here](https://blog.plasticlabs.ai/blog/A-Simple-Honcho-Primer).
|
||||
|
||||
|
|
@ -60,102 +58,136 @@ version of Honcho.
|
|||
The functionality of Honcho can be split into two different services: Storage
|
||||
and Insights.
|
||||
|
||||
|
||||
### Peer Paradigm
|
||||
|
||||
Honcho uses a peer-based model where both users and agents are represented as "peers". This unified approach enables:
|
||||
|
||||
- Multi-participant sessions with mixed human and AI agents
|
||||
- Configurable observation settings (which peers observe which others)
|
||||
- Flexible identity management for all participants
|
||||
- Support for complex multi-agent interactions
|
||||
|
||||
#### Key Features
|
||||
|
||||
- **Theory-of-Mind System**: Multiple implementation methods that extract facts from interactions and build comprehensive models of peer psychology
|
||||
- **Dialectic API**: Provides theory-of-mind informed responses that integrate long-term facts with current context
|
||||
- **Background Processing**: Asynchronous processing pipeline for expensive operations like representation updates and session summarization
|
||||
- **Multi-Provider Support**: Configurable LLM providers for different use cases
|
||||
|
||||
### Storage
|
||||
|
||||
Honcho contains several different primitives used for storing application and
|
||||
user data. This data is used for managing conversations, modeling user
|
||||
peer data. This data is used for managing conversations, modeling peer
|
||||
psychology, building RAG applications, and more.
|
||||
|
||||
The philosophy behind Honcho is to provide a platform that is user-centric and
|
||||
The philosophy behind Honcho is to provide a platform that is peer-centric and
|
||||
easily scalable from a single user to a million.
|
||||
|
||||
Below is a mapping of the different primitives.
|
||||
Below is a mapping of the different primitives and their relationships.
|
||||
|
||||
```
|
||||
Apps
|
||||
└── Users
|
||||
├── Sessions
|
||||
│ └── Messages
|
||||
├── Collections
|
||||
│ └── Documents
|
||||
└── Metamessages
|
||||
Workspaces
|
||||
├── Peers ←──────────────────┐
|
||||
│ ├── Sessions │
|
||||
│ ├── Collections │
|
||||
│ │ └── Documents │
|
||||
│ └── Messages (peer-level)│
|
||||
│ │
|
||||
└── Sessions ←───────────────┤ (many-to-many)
|
||||
├── Peers ───────────────┘
|
||||
└── Messages (session-level)
|
||||
```
|
||||
|
||||
**Relationship Details:**
|
||||
- A **Workspace** contains multiple **Peers**
|
||||
- **Peers** and **Sessions** have a many-to-many relationship (peers can participate in multiple sessions, sessions can have multiple peers)
|
||||
- **Messages** can exist at two levels:
|
||||
- **Session-level**: Communication between peers within a session
|
||||
- **Peer-level**: Data ingested by a peer to enhance its global representation
|
||||
- **Collections** belong to specific **Peers**
|
||||
- **Documents** are stored within **Collections**
|
||||
|
||||
Users familiar with APIs such as the OpenAI Assistants API will be familiar with
|
||||
much of the mapping here.
|
||||
|
||||
#### Apps
|
||||
#### Workspaces
|
||||
|
||||
This is the top level construct of Honcho. Developers can register different
|
||||
`Apps` for different assistants, agents, AI enabled features, etc. It is a way to
|
||||
isolation data between use cases.
|
||||
This is the top level construct of Honcho (formerly called Apps). Developers can register different
|
||||
`Workspaces` for different assistants, agents, AI enabled features, etc. It is a way to
|
||||
isolate data between use cases and provide multi-tenant capabilities.
|
||||
|
||||
**Users**
|
||||
#### Peers
|
||||
|
||||
Within an `App` everything revolves around a `User`. the `User` object
|
||||
literally represent a user of an application.
|
||||
Within a `Workspace` everything revolves around a `Peer`. The `Peer` object
|
||||
represents any participant in the system - whether human users or AI agents.
|
||||
This unified model enables complex multi-participant interactions.
|
||||
|
||||
#### Sessions
|
||||
|
||||
The `Session` object represents a set of interactions a `User` has with an
|
||||
`App`. Other application may refer to this as a thread or conversation.
|
||||
The `Session` object represents a set of interactions between `Peers` within a
|
||||
`Workspace`. Other applications may refer to this as a thread or conversation.
|
||||
Sessions can involve multiple peers with configurable observation settings.
|
||||
|
||||
**Messages**
|
||||
#### Messages
|
||||
|
||||
The `Message` represents an atomic interaction of a `User` in a `Session`.
|
||||
`Message`s are labeled as either a `User` or AI message.
|
||||
The `Message` represents an atomic data unit that can exist at two levels:
|
||||
|
||||
- **Session-level Messages**: Communication between peers within a session context
|
||||
- **Peer-level Messages**: Arbitrary data ingested by a peer to enhance its global representation (independent of any session)
|
||||
|
||||
All messages are labeled by their source peer and can be processed asynchronously to update theory-of-mind models. This flexible design allows for both conversational interactions and broader data ingestion for personality modeling.
|
||||
|
||||
#### Collections
|
||||
|
||||
At a high level a `Collection` is a named group of `Documents`. Developers
|
||||
familiar with RAG based applications will be familiar with these. `Collection`s
|
||||
familiar with RAG based applications will be familiar with these. `Collections`
|
||||
store vector embedded data that developers and agents can retrieve against using
|
||||
functions like cosine similarity.
|
||||
|
||||
Developers can create multiple `Collection`s for a user for different purposes
|
||||
Developers can create multiple `Collections` for a peer for different purposes
|
||||
such as modeling different personas, adding third-party data such as emails and
|
||||
PDF files, and more.
|
||||
PDF files, and more. Collections are also used internally by Honcho to store
|
||||
theory-of-mind representations.
|
||||
|
||||
#### Documents
|
||||
|
||||
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
|
||||
`Messages` and `Sessions` are created for a `User`, Honcho will asynchronously
|
||||
reason about the `User`'s psychology to derive facts about them and store them
|
||||
in a reserved `Collection`.
|
||||
`Messages` and `Sessions` are created for `Peers`, Honcho will asynchronously
|
||||
reason about peer psychology to derive facts about them and store them
|
||||
in reserved `Collections`.
|
||||
|
||||
The system uses a sophisticated message processing pipeline:
|
||||
|
||||
1. Messages are created via API
|
||||
2. Enqueued for background processing including:
|
||||
- `representation`: Update peer's theory of mind
|
||||
- `summary`: Create session summaries
|
||||
3. Session-based queue processing ensures proper ordering
|
||||
4. Results are stored internally in the vector database
|
||||
|
||||
To read more about how this works read our [Research Paper](https://arxiv.org/abs/2310.06983)
|
||||
|
||||
Developers can then leverage these insights in their application to better
|
||||
server `User` needs. The primary interface for using these insights is through
|
||||
serve peer needs. The primary interface for using these insights is through
|
||||
the [Dialectic Endpoint](https://blog.plasticlabs.ai/blog/Introducing-Honcho's-Dialectic-API).
|
||||
|
||||
This is a regular API endpoint that takes natural language requests to get data
|
||||
about the `User`. This robust design let's us use this single endpoint for all
|
||||
cases where extra personalization or information about the `User` is necessary.
|
||||
This is a regular API endpoint (`/peers/{peer_id}/chat`) that takes natural language requests to get data
|
||||
about the `Peer`. This robust design lets us use this single endpoint for all
|
||||
cases where extra personalization or information about the `Peer` is necessary.
|
||||
|
||||
A developer's application can treat Honcho as an oracle to the `User` and
|
||||
A developer's application can treat Honcho as an oracle to the `Peer` and
|
||||
consult it when necessary. Some examples of how to leverage the Dialectic
|
||||
API include:
|
||||
|
||||
- Asking Honcho for a theory-of-mind insight about the `User`
|
||||
- Asking Honcho to hydrate a prompt with data about the `User`s behavior
|
||||
- Asking Honcho for a 2nd opinion or approach about how to respond to the User
|
||||
- Asking Honcho for a theory-of-mind insight about the `Peer`
|
||||
- Asking Honcho to hydrate a prompt with data about the `Peer`s behavior
|
||||
- Asking Honcho for a 2nd opinion or approach about how to respond to the Peer
|
||||
- Getting personalized responses that incorporate long-term facts and context
|
||||
|
||||
## License
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,14 @@ def get_schema() -> str:
|
|||
return getenv("DATABASE_SCHEMA", "public")
|
||||
|
||||
|
||||
def table_exists(table_name: str, inspector: Optional[sa.Inspector] = None) -> bool:
|
||||
"""Check if a table exists in the database."""
|
||||
if inspector is None:
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
schema = get_schema()
|
||||
return inspector.has_table(table_name, schema=schema)
|
||||
|
||||
|
||||
def column_exists(
|
||||
table_name: str, column_name: str, inspector: Optional[sa.Inspector] = None
|
||||
) -> bool:
|
||||
|
|
@ -43,3 +51,24 @@ def index_exists(
|
|||
indexes = inspector.get_indexes(table_name, schema=schema)
|
||||
return any(idx["name"] == index_name for idx in indexes)
|
||||
|
||||
|
||||
def constraint_exists(
|
||||
table_name: str,
|
||||
constraint_name: str,
|
||||
type: str,
|
||||
inspector: Optional[sa.Inspector] = None,
|
||||
) -> bool:
|
||||
"""Check if a constraint exists in a table."""
|
||||
if inspector is None:
|
||||
inspector = sa.inspect(op.get_bind())
|
||||
schema = get_schema()
|
||||
if type == "check":
|
||||
constraints = inspector.get_check_constraints(table_name, schema=schema)
|
||||
elif type == "unique":
|
||||
constraints = inspector.get_unique_constraints(table_name, schema=schema)
|
||||
elif type == "primary":
|
||||
constraint = inspector.get_pk_constraint(table_name, schema=schema)
|
||||
return constraint["name"] == constraint_name
|
||||
else:
|
||||
raise ValueError(f"Invalid constraint type: {type}")
|
||||
return any(constraint["name"] == constraint_name for constraint in constraints)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "honcho"
|
||||
version = "1.1.0"
|
||||
version = "2.0.0"
|
||||
description = "Honcho Server"
|
||||
authors = [
|
||||
{name = "Plastic Labs", email = "hello@plasticlabs.ai"},
|
||||
|
|
@ -25,6 +25,7 @@ dependencies = [
|
|||
"langfuse>=2.57.1",
|
||||
"pyjwt>=2.10.0",
|
||||
"google-genai>=1.10.0",
|
||||
"tiktoken>=0.9.0",
|
||||
]
|
||||
[tool.uv]
|
||||
dev-dependencies = [
|
||||
|
|
|
|||
174
src/agent.py
174
src/agent.py
|
|
@ -2,7 +2,6 @@ import asyncio
|
|||
import json
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import Any, Optional
|
||||
|
||||
import sentry_sdk
|
||||
|
|
@ -24,8 +23,6 @@ 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"
|
||||
|
||||
|
|
@ -131,9 +128,9 @@ class Dialectic:
|
|||
|
||||
@observe()
|
||||
async def chat(
|
||||
app_id: str,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
workspace_name: str,
|
||||
peer_name: str,
|
||||
session_name: str | None,
|
||||
queries: str | list[str],
|
||||
stream: bool = False,
|
||||
) -> schemas.DialecticResponse | MessageStreamManager:
|
||||
|
|
@ -148,58 +145,77 @@ async def chat(
|
|||
3. Combines both into a fresh user representation
|
||||
4. Uses this representation to answer the query
|
||||
5. Saves the representation for future use
|
||||
|
||||
Args:
|
||||
workspace_name: The workspace name
|
||||
peer_name: The peer name
|
||||
session_name: The session name. If None, this queries the global representation.
|
||||
queries: The queries to ask the Dialectic API
|
||||
stream: Whether to stream the response
|
||||
|
||||
Returns:
|
||||
Either a string or a stream of messages from the LLM provider, depending on the stream flag
|
||||
"""
|
||||
# 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(f"Received query: {final_query} for session {session_name}")
|
||||
logger.debug("Starting on-demand user representation generation")
|
||||
|
||||
start_time = asyncio.get_event_loop().time()
|
||||
|
||||
# Setup phase - create resources we'll need for all operations
|
||||
|
||||
# 1. Fetch latest user message & chat history
|
||||
# 1. Fetch latest peer message & chat history
|
||||
async with tracked_db("chat.load_history") as db_history:
|
||||
stmt = (
|
||||
select(models.Message)
|
||||
.where(models.Message.app_id == app_id)
|
||||
.where(models.Message.user_id == user_id)
|
||||
.where(models.Message.session_id == session_id)
|
||||
.where(models.Message.is_user)
|
||||
.where(models.Message.workspace_name == workspace_name)
|
||||
.where(models.Message.peer_name == peer_name)
|
||||
.order_by(models.Message.id.desc())
|
||||
.limit(1)
|
||||
)
|
||||
if session_name:
|
||||
stmt = stmt.where(models.Message.session_name == session_name)
|
||||
latest_messages = await db_history.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}")
|
||||
|
||||
chat_history, _, _ = await history.get_summarized_history(
|
||||
db_history, 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")
|
||||
if session_name:
|
||||
chat_history = await history.get_summarized_history(
|
||||
db_history,
|
||||
workspace_name,
|
||||
session_name,
|
||||
peer_name,
|
||||
summary_type=history.SummaryType.SHORT,
|
||||
)
|
||||
if not chat_history:
|
||||
logger.warning(f"No chat history found for session {session_name}")
|
||||
chat_history = (
|
||||
f"someone asked this about the user's message: {final_query}"
|
||||
)
|
||||
logger.debug(
|
||||
f"Workspace: {workspace_name}, Peer: {peer_name}, Session: {session_name}"
|
||||
)
|
||||
else:
|
||||
chat_history = ""
|
||||
logger.debug("Retrieved chat history: %s lines", len(chat_history.split("\n")))
|
||||
|
||||
# Run short-term inference and long-term facts in parallel
|
||||
async def fetch_long_term():
|
||||
async with tracked_db("chat.get_collection") as db_embed:
|
||||
collection = await crud.get_or_create_user_protected_collection(
|
||||
db_embed, app_id, user_id
|
||||
name = "global_representation" if session_name is None else ""
|
||||
collection = await crud.get_or_create_collection(
|
||||
db_embed, workspace_name, peer_name, collection_name=name
|
||||
)
|
||||
collection_id = (
|
||||
collection.public_id
|
||||
) # Extract the ID while session is active
|
||||
facts = await get_long_term_facts(final_query, app_id, user_id, collection_id)
|
||||
collection_name = collection.name # Extract the ID while session is active
|
||||
facts = await get_long_term_facts(
|
||||
final_query, workspace_name, peer_name, collection_name
|
||||
)
|
||||
return facts
|
||||
|
||||
long_term_task = asyncio.create_task(fetch_long_term())
|
||||
short_term_task = asyncio.create_task(run_tom_inference(chat_history, session_id))
|
||||
short_term_task = asyncio.create_task(run_tom_inference(chat_history))
|
||||
|
||||
facts, tom_inference = await asyncio.gather(long_term_task, short_term_task)
|
||||
logger.debug(f"Retrieved {len(facts)} facts from long-term memory")
|
||||
|
|
@ -209,9 +225,9 @@ async def chat(
|
|||
logger.debug("Generating user representation")
|
||||
async with tracked_db("chat.generate_user_representation") as db_rep:
|
||||
user_representation = await generate_user_representation(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
workspace_name,
|
||||
peer_name,
|
||||
session_name,
|
||||
chat_history=chat_history,
|
||||
tom_inference=tom_inference,
|
||||
facts=facts,
|
||||
|
|
@ -234,8 +250,8 @@ async def chat(
|
|||
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,
|
||||
session_id=session_name,
|
||||
user_id=peer_name,
|
||||
release=os.getenv("SENTRY_RELEASE"),
|
||||
metadata={"environment": os.getenv("SENTRY_ENVIRONMENT")},
|
||||
)
|
||||
|
|
@ -260,14 +276,19 @@ async def chat(
|
|||
|
||||
|
||||
async def get_long_term_facts(
|
||||
query: str, app_id: str, user_id: str, collection_id: str
|
||||
query: str,
|
||||
workspace_name: str,
|
||||
peer_name: str,
|
||||
collection_name: str,
|
||||
) -> 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
|
||||
workspace_name: The workspace name
|
||||
peer_name: The peer name
|
||||
collection_name: The collection name
|
||||
|
||||
Returns:
|
||||
List of retrieved facts
|
||||
|
|
@ -285,9 +306,9 @@ async def get_long_term_facts(
|
|||
logger.debug(f"Starting query {i + 1}/{len(search_queries)}: {search_query}")
|
||||
query_start = asyncio.get_event_loop().time()
|
||||
query_embedding_store = CollectionEmbeddingStore(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
workspace_name=workspace_name,
|
||||
peer_name=peer_name,
|
||||
collection_name=collection_name,
|
||||
)
|
||||
facts = await query_embedding_store.get_relevant_facts(
|
||||
search_query, top_k=10, max_distance=0.85
|
||||
|
|
@ -314,24 +335,23 @@ async def get_long_term_facts(
|
|||
return list(retrieved_facts)
|
||||
|
||||
|
||||
async def run_tom_inference(chat_history: str, session_id: str) -> str:
|
||||
async def run_tom_inference(chat_history: 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}")
|
||||
logger.debug("Running ToM inference")
|
||||
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=""
|
||||
chat_history, user_representation="", method="single_prompt"
|
||||
)
|
||||
|
||||
# Extract the prediction from the response
|
||||
|
|
@ -407,9 +427,9 @@ async def generate_semantic_queries(query: str) -> list[str]:
|
|||
|
||||
|
||||
async def generate_user_representation(
|
||||
app_id: str,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
workspace_name: str,
|
||||
peer_name: str,
|
||||
session_name: str | None,
|
||||
chat_history: str,
|
||||
tom_inference: str,
|
||||
facts: list[str],
|
||||
|
|
@ -419,8 +439,9 @@ async def generate_user_representation(
|
|||
) -> 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.
|
||||
Save it to peer metadata if no session is provided (global-level), or save it to
|
||||
session-peers table metadata if a session is provided (local-level).
|
||||
If session-level, uses existing representations from the same session for continuity.
|
||||
|
||||
Returns:
|
||||
The generated user representation.
|
||||
|
|
@ -429,24 +450,10 @@ async def generate_user_representation(
|
|||
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)
|
||||
.where(
|
||||
models.Metamessage.session_id == session_id
|
||||
) # only from the same session
|
||||
.where(models.Metamessage.label == 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."
|
||||
latest_representation = await crud.get_working_representation(
|
||||
db, workspace_name, peer_name, session_name
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"Found previous representation: {len(latest_representation)} characters"
|
||||
)
|
||||
|
|
@ -457,7 +464,6 @@ async def generate_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,
|
||||
user_representation=latest_representation,
|
||||
tom_inference=tom_inference,
|
||||
|
|
@ -479,36 +485,24 @@ 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 message_id is provided, save the representation as metadata
|
||||
if not representation:
|
||||
logger.debug("Empty representation, skipping save")
|
||||
elif not message_id:
|
||||
logger.debug("No message_id, skipping save")
|
||||
else:
|
||||
logger.debug(f"Saving representation to message_id: {message_id}")
|
||||
save_start = asyncio.get_event_loop().time()
|
||||
try:
|
||||
# First check if message exists
|
||||
message_check_stmt = select(models.Message).where(
|
||||
models.Message.public_id == message_id
|
||||
await crud.set_working_representation(
|
||||
db,
|
||||
representation,
|
||||
workspace_name,
|
||||
peer_name,
|
||||
session_name,
|
||||
)
|
||||
message_check = await 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(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
message_id=message_id if message_id else None,
|
||||
label=USER_REPRESENTATION_METAMESSAGE_TYPE,
|
||||
content=representation,
|
||||
h_metadata={},
|
||||
)
|
||||
db.add(metamessage)
|
||||
await db.commit()
|
||||
save_time = asyncio.get_event_loop().time() - save_start
|
||||
logger.debug(f"Representation saved in {save_time:.2f}s")
|
||||
save_time = asyncio.get_event_loop().time() - save_start
|
||||
logger.debug(f"Representation saved in {save_time:.2f}s")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during save DB operation: {str(e)}")
|
||||
await db.rollback()
|
||||
|
|
|
|||
2223
src/crud.py
2223
src/crud.py
File diff suppressed because it is too large
Load Diff
|
|
@ -1,5 +1,6 @@
|
|||
import contextvars
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from sqlalchemy import MetaData, create_engine, text
|
||||
|
|
@ -11,7 +12,9 @@ load_dotenv()
|
|||
connect_args = {"prepare_threshold": None}
|
||||
|
||||
# Context variable to store request context
|
||||
request_context = contextvars.ContextVar("request_context", default=None)
|
||||
request_context: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
|
||||
"request_context", default=None
|
||||
)
|
||||
|
||||
engine = create_async_engine(
|
||||
os.environ["CONNECTION_URI"],
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import logging
|
||||
import os
|
||||
from typing import Literal
|
||||
|
||||
import sentry_sdk
|
||||
from langfuse.decorators import observe
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from rich.console import Console
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
|
|
@ -20,70 +22,82 @@ TOM_METHOD = os.getenv("TOM_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)
|
||||
class PayloadSchema(BaseModel):
|
||||
"""
|
||||
Schema for validating payload data in process_item function.
|
||||
Ensures all required fields are present with correct types and prevents injection risks.
|
||||
"""
|
||||
|
||||
content: str
|
||||
workspace_name: str
|
||||
sender_name: str
|
||||
target_name: str
|
||||
session_name: str | None
|
||||
message_id: int
|
||||
task_type: Literal["representation", "summary"]
|
||||
|
||||
class Config:
|
||||
# Forbid extra fields to prevent injection of unexpected data
|
||||
extra = "forbid"
|
||||
|
||||
|
||||
async def process_item(db: AsyncSession, payload: dict):
|
||||
# Validate payload structure and types before processing
|
||||
try:
|
||||
validated_payload = PayloadSchema(**payload)
|
||||
except ValidationError as e:
|
||||
logger.error("Invalid payload received: %s. Payload: %s", str(e), payload)
|
||||
raise ValueError(f"Invalid payload structure: {str(e)}") from e
|
||||
|
||||
logger.debug(
|
||||
f"process_item received payload: {payload['message_id']} is_user={payload['is_user']}"
|
||||
"process_item received payload for message %s in session %s",
|
||||
validated_payload.message_id,
|
||||
validated_payload.session_name,
|
||||
)
|
||||
|
||||
processing_args = [
|
||||
payload["content"],
|
||||
payload["app_id"],
|
||||
payload["user_id"],
|
||||
payload["session_id"],
|
||||
payload["message_id"],
|
||||
validated_payload.content,
|
||||
validated_payload.workspace_name,
|
||||
validated_payload.sender_name,
|
||||
validated_payload.target_name,
|
||||
validated_payload.session_name,
|
||||
validated_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']}")
|
||||
if validated_payload.task_type == "representation":
|
||||
logger.debug(
|
||||
"Processing message %s in %s",
|
||||
validated_payload.message_id,
|
||||
validated_payload.session_name,
|
||||
)
|
||||
await process_message(*processing_args)
|
||||
logger.debug(
|
||||
"Finished processing message %s in %s %s",
|
||||
validated_payload.message_id,
|
||||
"session" if validated_payload.session_name else "peer",
|
||||
validated_payload.session_name
|
||||
if validated_payload.session_name
|
||||
else validated_payload.sender_name,
|
||||
)
|
||||
await summarize_if_needed(
|
||||
db,
|
||||
payload["app_id"],
|
||||
payload["session_id"],
|
||||
payload["user_id"],
|
||||
payload["message_id"],
|
||||
validated_payload.workspace_name,
|
||||
validated_payload.session_name,
|
||||
validated_payload.sender_name,
|
||||
validated_payload.message_id,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
@sentry_sdk.trace
|
||||
# @observe()
|
||||
async def process_ai_message(
|
||||
content: str,
|
||||
app_id: str,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
message_id: str,
|
||||
db: AsyncSession,
|
||||
):
|
||||
"""
|
||||
Process an AI message. Make a prediction about what the user is going to say to it.
|
||||
"""
|
||||
console.print(f"Processing AI message: {content}", style="bright_magenta")
|
||||
|
||||
|
||||
@sentry_sdk.trace
|
||||
@observe()
|
||||
async def process_user_message(
|
||||
async def process_message(
|
||||
content: str,
|
||||
app_id: str,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
message_id: str,
|
||||
workspace_name: str,
|
||||
peer_name: str,
|
||||
target_name: str,
|
||||
session_name: str | None,
|
||||
message_id: int,
|
||||
db: AsyncSession,
|
||||
):
|
||||
"""
|
||||
|
|
@ -92,18 +106,32 @@ async def process_user_message(
|
|||
"""
|
||||
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
|
||||
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
|
||||
logger.debug(
|
||||
"Starting fact extraction for user message %s in %s %s",
|
||||
message_id,
|
||||
"session" if session_name else "peer",
|
||||
session_name if session_name else peer_name,
|
||||
)
|
||||
chat_history_str = f"{short_history_text}\nhuman: {content}"
|
||||
|
||||
if session_name:
|
||||
# Get chat history and append current message
|
||||
logger.debug(
|
||||
"Retrieving chat history for %s %s",
|
||||
"session" if session_name else "peer",
|
||||
session_name if session_name else peer_name,
|
||||
)
|
||||
short_history_text = await history.get_summarized_history(
|
||||
db,
|
||||
workspace_name,
|
||||
session_name,
|
||||
peer_name,
|
||||
cutoff=message_id,
|
||||
summary_type=history.SummaryType.SHORT,
|
||||
)
|
||||
|
||||
chat_history_str = f"{short_history_text}\nuser: {content}"
|
||||
else:
|
||||
chat_history_str = f"user: {content}"
|
||||
|
||||
# Extract facts from chat history
|
||||
logger.debug("Extracting facts from chat history")
|
||||
|
|
@ -114,14 +142,21 @@ async def process_user_message(
|
|||
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
|
||||
logger.debug(
|
||||
f"Setting up embedding store for workspace: {workspace_name}, peer: {peer_name}"
|
||||
)
|
||||
collection_name = (
|
||||
crud.construct_collection_name(peer_name, target_name)
|
||||
if peer_name != target_name
|
||||
else "global_representation"
|
||||
)
|
||||
collection = await crud.get_or_create_collection(
|
||||
db, workspace_name, peer_name, collection_name
|
||||
)
|
||||
embedding_store = CollectionEmbeddingStore(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection.public_id, # type: ignore
|
||||
workspace_name=workspace_name,
|
||||
peer_name=peer_name,
|
||||
collection_name=collection.name,
|
||||
)
|
||||
|
||||
# Filter out facts that are duplicates of existing facts in the vector store
|
||||
|
|
@ -150,18 +185,30 @@ async def process_user_message(
|
|||
|
||||
|
||||
async def summarize_if_needed(
|
||||
db: AsyncSession, app_id: str, session_id: str, user_id: str, message_id: str
|
||||
db: AsyncSession,
|
||||
workspace_name: str,
|
||||
session_name: str | None,
|
||||
peer_name: str,
|
||||
message_id: int,
|
||||
):
|
||||
if not session_name:
|
||||
return
|
||||
|
||||
summary_start = os.times()[4]
|
||||
logger.debug("Checking if summaries should be created")
|
||||
logger.debug("Checking if summaries should be created for session %s", session_name)
|
||||
|
||||
# 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
|
||||
db,
|
||||
workspace_name,
|
||||
session_name,
|
||||
peer_name,
|
||||
message_id,
|
||||
summary_type=history.SummaryType.SHORT,
|
||||
)
|
||||
|
||||
if should_create_short:
|
||||
|
|
@ -173,7 +220,12 @@ async def summarize_if_needed(
|
|||
long_messages,
|
||||
latest_long_summary,
|
||||
) = await history.should_create_summary(
|
||||
db, session_id, summary_type=history.SummaryType.LONG
|
||||
db,
|
||||
workspace_name,
|
||||
session_name,
|
||||
peer_name,
|
||||
message_id,
|
||||
summary_type=history.SummaryType.LONG,
|
||||
)
|
||||
|
||||
# STEP 3: If we need a long summary, create it first before creating the short summary
|
||||
|
|
@ -183,26 +235,23 @@ async def summarize_if_needed(
|
|||
)
|
||||
try:
|
||||
# Get previous long summary context if available
|
||||
previous_long_summary = (
|
||||
latest_long_summary.content if latest_long_summary else None
|
||||
previous_long_summary_text = (
|
||||
latest_long_summary["content"] if latest_long_summary else None
|
||||
)
|
||||
|
||||
# Create a new long summary
|
||||
long_summary_text = await history.create_summary(
|
||||
new_long_summary = await history.create_summary(
|
||||
messages=long_messages,
|
||||
previous_summary=previous_long_summary,
|
||||
previous_summary_text=previous_long_summary_text,
|
||||
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,
|
||||
app_id=app_id,
|
||||
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,
|
||||
|
||||
# Save the long summary
|
||||
await history.save_summary(
|
||||
db,
|
||||
new_long_summary,
|
||||
workspace_name,
|
||||
session_name,
|
||||
)
|
||||
logger.debug("Long summary created and saved successfully")
|
||||
except Exception as e:
|
||||
|
|
@ -217,25 +266,23 @@ async def summarize_if_needed(
|
|||
f"Creating new short summary covering {len(short_messages)} messages"
|
||||
)
|
||||
try:
|
||||
previous_summary = (
|
||||
latest_long_summary.content if latest_long_summary else None
|
||||
previous_long_summary_text = (
|
||||
latest_long_summary["content"] if latest_long_summary else None
|
||||
)
|
||||
|
||||
# Create a new short summary
|
||||
short_summary_text = await history.create_summary(
|
||||
new_short_summary = await history.create_summary(
|
||||
messages=short_messages,
|
||||
previous_summary=previous_summary,
|
||||
previous_summary_text=previous_long_summary_text,
|
||||
summary_type=history.SummaryType.SHORT,
|
||||
)
|
||||
# Save the short summary as a metamessage
|
||||
await history.save_summary_metamessage(
|
||||
db=db,
|
||||
app_id=app_id,
|
||||
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,
|
||||
|
||||
# Save the short summary
|
||||
await history.save_summary(
|
||||
db,
|
||||
new_short_summary,
|
||||
workspace_name,
|
||||
session_name,
|
||||
)
|
||||
logger.debug("Short summary created and saved successfully")
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import asyncio
|
||||
import os
|
||||
import signal
|
||||
from datetime import datetime, timedelta
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from logging import getLogger
|
||||
|
||||
import sentry_sdk
|
||||
|
|
@ -21,11 +23,30 @@ logger = getLogger(__name__)
|
|||
load_dotenv()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkUnit:
|
||||
"""
|
||||
Represents a unit of work in the queue system.
|
||||
|
||||
A work unit is uniquely identified by the combination of session_id,
|
||||
sender_name, target_name, and task_type. This allows multiple workers
|
||||
to process different work units from the same session in parallel.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
sender_name: str
|
||||
target_name: str
|
||||
task_type: str
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"({self.session_id}, {self.sender_name}, {self.target_name}, {self.task_type})"
|
||||
|
||||
|
||||
class QueueManager:
|
||||
def __init__(self):
|
||||
self.shutdown_event = asyncio.Event()
|
||||
self.active_tasks: set[asyncio.Task] = set()
|
||||
self.owned_sessions: set[int] = set()
|
||||
self.owned_work_units: set[WorkUnit] = set()
|
||||
self.queue_empty_flag = asyncio.Event()
|
||||
|
||||
# Initialize from environment
|
||||
|
|
@ -47,13 +68,13 @@ class QueueManager:
|
|||
self.active_tasks.add(task)
|
||||
task.add_done_callback(self.active_tasks.discard)
|
||||
|
||||
def track_session(self, session_id: int):
|
||||
"""Track a new session owned by this process"""
|
||||
self.owned_sessions.add(session_id)
|
||||
def track_work_unit(self, work_unit: WorkUnit):
|
||||
"""Track a new work unit owned by this process"""
|
||||
self.owned_work_units.add(work_unit)
|
||||
|
||||
def untrack_session(self, session_id: int):
|
||||
"""Remove a session from tracking"""
|
||||
self.owned_sessions.discard(session_id)
|
||||
def untrack_work_unit(self, work_unit: WorkUnit):
|
||||
"""Remove a work unit from tracking"""
|
||||
self.owned_work_units.discard(work_unit)
|
||||
|
||||
async def initialize(self):
|
||||
"""Setup signal handlers and start the main polling loop"""
|
||||
|
|
@ -87,19 +108,25 @@ class QueueManager:
|
|||
await asyncio.gather(*self.active_tasks, return_exceptions=True)
|
||||
|
||||
async def cleanup(self):
|
||||
"""Clean up owned sessions"""
|
||||
if self.owned_sessions:
|
||||
logger.info(f"Cleaning up {len(self.owned_sessions)} owned sessions...")
|
||||
"""Clean up owned work units"""
|
||||
if self.owned_work_units:
|
||||
logger.info(f"Cleaning up {len(self.owned_work_units)} owned work units...")
|
||||
try:
|
||||
# Use the tracked_db dependency for transaction safety
|
||||
async with tracked_db("queue_cleanup") as db:
|
||||
await db.execute(
|
||||
delete(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.session_id.in_(
|
||||
self.owned_sessions
|
||||
for work_unit in self.owned_work_units:
|
||||
await db.execute(
|
||||
delete(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.session_id
|
||||
== work_unit.session_id,
|
||||
models.ActiveQueueSession.sender_name
|
||||
== work_unit.sender_name,
|
||||
models.ActiveQueueSession.target_name
|
||||
== work_unit.target_name,
|
||||
models.ActiveQueueSession.task_type
|
||||
== work_unit.task_type,
|
||||
)
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
logger.info("Cleanup completed successfully")
|
||||
except Exception as e:
|
||||
|
|
@ -111,34 +138,66 @@ class QueueManager:
|
|||
# Polling and Scheduling #
|
||||
##########################
|
||||
|
||||
async def get_available_sessions(self, db: AsyncSession):
|
||||
"""Get available sessions that aren't being processed"""
|
||||
# Clean up stale sessions
|
||||
five_minutes_ago = datetime.utcnow() - timedelta(minutes=5)
|
||||
async def get_available_work_units(self, db: AsyncSession) -> Sequence[WorkUnit]:
|
||||
"""
|
||||
Get available work units that aren't being processed.
|
||||
Returns a list of WorkUnit objects.
|
||||
"""
|
||||
# Clean up stale work units
|
||||
five_minutes_ago = datetime.now(UTC) - timedelta(minutes=5)
|
||||
await db.execute(
|
||||
delete(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.last_updated < five_minutes_ago
|
||||
)
|
||||
)
|
||||
|
||||
# Get available sessions
|
||||
# Create the JSON path expressions once to ensure they're identical in SELECT and GROUP BY
|
||||
sender_name_expr = models.QueueItem.payload["sender_name"].astext
|
||||
target_name_expr = models.QueueItem.payload["target_name"].astext
|
||||
task_type_expr = models.QueueItem.payload["task_type"].astext
|
||||
|
||||
# Get available work units by extracting sender_name, target_name, task_type from payload
|
||||
# We need to join with ActiveQueueSession to find units that aren't already being processed
|
||||
result = await db.execute(
|
||||
select(models.QueueItem.session_id)
|
||||
select(
|
||||
models.QueueItem.session_id,
|
||||
sender_name_expr.label("sender_name"),
|
||||
target_name_expr.label("target_name"),
|
||||
task_type_expr.label("task_type"),
|
||||
)
|
||||
.outerjoin(
|
||||
models.ActiveQueueSession,
|
||||
models.QueueItem.session_id == models.ActiveQueueSession.session_id,
|
||||
(models.QueueItem.session_id == models.ActiveQueueSession.session_id)
|
||||
& (sender_name_expr == models.ActiveQueueSession.sender_name)
|
||||
& (target_name_expr == models.ActiveQueueSession.target_name)
|
||||
& (task_type_expr == models.ActiveQueueSession.task_type),
|
||||
)
|
||||
.where(models.QueueItem.processed == False)
|
||||
.where(~models.QueueItem.processed)
|
||||
.where(
|
||||
models.ActiveQueueSession.session_id == None
|
||||
) # Only sessions not in active_sessions
|
||||
.group_by(models.QueueItem.session_id)
|
||||
.limit(1)
|
||||
models.ActiveQueueSession.id == None # noqa: E711
|
||||
) # Only work units not in active_queue_sessions
|
||||
.group_by(
|
||||
models.QueueItem.session_id,
|
||||
sender_name_expr,
|
||||
target_name_expr,
|
||||
task_type_expr,
|
||||
)
|
||||
.limit(self.workers) # Process multiple work units in parallel
|
||||
)
|
||||
return result.scalars().all()
|
||||
|
||||
rows = result.fetchall()
|
||||
return [
|
||||
WorkUnit(
|
||||
session_id=row.session_id,
|
||||
sender_name=row.sender_name,
|
||||
target_name=row.target_name,
|
||||
task_type=row.task_type,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
async def polling_loop(self):
|
||||
"""Main polling loop to find and process new sessions"""
|
||||
"""Main polling loop to find and process new work units"""
|
||||
logger.debug("Starting polling loop")
|
||||
try:
|
||||
while not self.shutdown_event.is_set():
|
||||
|
|
@ -157,35 +216,38 @@ class QueueManager:
|
|||
# Use the dependency for transaction safety
|
||||
async with tracked_db("queue_polling_loop") as db:
|
||||
try:
|
||||
new_sessions = await self.get_available_sessions(db)
|
||||
new_work_units = await self.get_available_work_units(db)
|
||||
|
||||
if new_sessions and not self.shutdown_event.is_set():
|
||||
for session_id in new_sessions:
|
||||
if new_work_units and not self.shutdown_event.is_set():
|
||||
for work_unit in new_work_units:
|
||||
try:
|
||||
# Try to claim the session
|
||||
# Try to claim the work unit
|
||||
await db.execute(
|
||||
insert(models.ActiveQueueSession).values(
|
||||
session_id=session_id,
|
||||
session_id=work_unit.session_id,
|
||||
sender_name=work_unit.sender_name,
|
||||
target_name=work_unit.target_name,
|
||||
task_type=work_unit.task_type,
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
# Track this session
|
||||
self.track_session(session_id)
|
||||
# Track this work unit
|
||||
self.track_work_unit(work_unit)
|
||||
logger.debug(
|
||||
f"Claimed session {session_id} for processing"
|
||||
f"Claimed work unit {work_unit} for processing"
|
||||
)
|
||||
|
||||
# Create a new task for processing this session
|
||||
# Create a new task for processing this work unit
|
||||
if not self.shutdown_event.is_set():
|
||||
task = asyncio.create_task(
|
||||
self.process_session(session_id)
|
||||
self.process_work_unit(work_unit)
|
||||
)
|
||||
self.add_task(task)
|
||||
except IntegrityError:
|
||||
# Note: rollback is handled by tracked_db dependency
|
||||
logger.debug(
|
||||
f"Failed to claim session {session_id}, already owned"
|
||||
f"Failed to claim work unit {work_unit}, already owned"
|
||||
)
|
||||
else:
|
||||
self.queue_empty_flag.set()
|
||||
|
|
@ -204,29 +266,29 @@ 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 def process_work_unit(self, work_unit: WorkUnit):
|
||||
"""Process all messages for a specific work unit"""
|
||||
logger.debug(f"Starting to process work unit {work_unit}")
|
||||
# Use the tracked_db dependency for transaction safety
|
||||
async with (
|
||||
self.semaphore,
|
||||
tracked_db("queue_process_session") as db,
|
||||
): # Hold the semaphore for the entire session duration
|
||||
tracked_db("queue_process_work_unit") as db,
|
||||
): # Hold the semaphore for the entire work unit duration
|
||||
try:
|
||||
message_count = 0
|
||||
while not self.shutdown_event.is_set():
|
||||
message = await self.get_next_message(db, session_id)
|
||||
message = await self.get_next_message(db, work_unit)
|
||||
if not message:
|
||||
logger.debug(f"No more messages for session {session_id}")
|
||||
logger.debug(f"No more messages for work unit {work_unit}")
|
||||
break
|
||||
|
||||
message_count += 1
|
||||
logger.debug(
|
||||
f"Processing message {message.id} for session {session_id} (message {message_count})"
|
||||
f"Processing message {message.id} for work unit {work_unit} (message {message_count})"
|
||||
)
|
||||
try:
|
||||
logger.info(
|
||||
f"Processing message {message.id} from session {session_id}"
|
||||
f"Processing message {message.id} from work unit {work_unit}"
|
||||
)
|
||||
await process_item(db, payload=message.payload)
|
||||
logger.debug(f"Successfully processed message {message.id}")
|
||||
|
|
@ -245,39 +307,57 @@ class QueueManager:
|
|||
|
||||
if self.shutdown_event.is_set():
|
||||
logger.debug(
|
||||
f"Shutdown requested, stopping processing for session {session_id}"
|
||||
f"Shutdown requested, stopping processing for work unit {work_unit}"
|
||||
)
|
||||
break
|
||||
|
||||
# Update last_updated timestamp to show this session is still being processed
|
||||
# Update last_updated timestamp to show this work unit is still being processed
|
||||
await db.execute(
|
||||
update(models.ActiveQueueSession)
|
||||
.where(models.ActiveQueueSession.session_id == session_id)
|
||||
.where(
|
||||
models.ActiveQueueSession.session_id
|
||||
== work_unit.session_id,
|
||||
models.ActiveQueueSession.sender_name
|
||||
== work_unit.sender_name,
|
||||
models.ActiveQueueSession.target_name
|
||||
== work_unit.target_name,
|
||||
models.ActiveQueueSession.task_type == work_unit.task_type,
|
||||
)
|
||||
.values(last_updated=func.now())
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
logger.debug(
|
||||
f"Completed processing session {session_id}, processed {message_count} messages"
|
||||
f"Completed processing work unit {work_unit}, processed {message_count} messages"
|
||||
)
|
||||
finally:
|
||||
# Remove session from active_sessions when done
|
||||
logger.debug(f"Removing session {session_id} from active sessions")
|
||||
# Remove work unit from active_queue_sessions when done
|
||||
logger.debug(f"Removing work unit {work_unit} from active sessions")
|
||||
await db.execute(
|
||||
delete(models.ActiveQueueSession).where(
|
||||
models.ActiveQueueSession.session_id == session_id
|
||||
models.ActiveQueueSession.session_id == work_unit.session_id,
|
||||
models.ActiveQueueSession.sender_name == work_unit.sender_name,
|
||||
models.ActiveQueueSession.target_name == work_unit.target_name,
|
||||
models.ActiveQueueSession.task_type == work_unit.task_type,
|
||||
)
|
||||
)
|
||||
await db.commit()
|
||||
self.untrack_session(session_id)
|
||||
self.untrack_work_unit(work_unit)
|
||||
|
||||
@sentry_sdk.trace
|
||||
async def get_next_message(self, db: AsyncSession, session_id: int):
|
||||
"""Get the next unprocessed message for a session"""
|
||||
async def get_next_message(self, db: AsyncSession, work_unit: WorkUnit):
|
||||
"""Get the next unprocessed message for a specific work unit"""
|
||||
result = await db.execute(
|
||||
select(models.QueueItem)
|
||||
.where(models.QueueItem.session_id == session_id)
|
||||
.where(models.QueueItem.processed == False)
|
||||
.where(models.QueueItem.session_id == work_unit.session_id)
|
||||
.where(
|
||||
models.QueueItem.payload["sender_name"].astext == work_unit.sender_name
|
||||
)
|
||||
.where(
|
||||
models.QueueItem.payload["target_name"].astext == work_unit.target_name
|
||||
)
|
||||
.where(models.QueueItem.payload["task_type"].astext == work_unit.task_type)
|
||||
.where(~models.QueueItem.processed)
|
||||
.order_by(models.QueueItem.id)
|
||||
.with_for_update(skip_locked=True)
|
||||
.limit(1)
|
||||
|
|
|
|||
|
|
@ -11,18 +11,17 @@ from .single_prompt import (
|
|||
|
||||
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
|
||||
chat_history, user_representation, **kwargs
|
||||
)
|
||||
elif method == "single_prompt":
|
||||
return await get_tom_inference_single_prompt(
|
||||
chat_history, session_id, user_representation, **kwargs
|
||||
chat_history, user_representation, **kwargs
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
|
@ -30,7 +29,6 @@ async def get_tom_inference(
|
|||
|
||||
async def get_user_representation(
|
||||
chat_history: str,
|
||||
session_id: str,
|
||||
user_representation: str = "None",
|
||||
tom_inference: str = "None",
|
||||
method: str = "conversational",
|
||||
|
|
@ -38,15 +36,15 @@ async def get_user_representation(
|
|||
) -> str:
|
||||
if method == "conversational":
|
||||
return await get_user_representation_conversational(
|
||||
chat_history, session_id, user_representation, tom_inference, **kwargs
|
||||
chat_history, 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
|
||||
chat_history, 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
|
||||
chat_history, user_representation, tom_inference, **kwargs
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
|
|
|||
|
|
@ -2,11 +2,10 @@ import os
|
|||
|
||||
import sentry_sdk
|
||||
from anthropic import Anthropic
|
||||
from anthropic.types import MessageParam, TextBlock
|
||||
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
|
||||
|
||||
# Initialize the Anthropic client
|
||||
anthropic = Anthropic(
|
||||
api_key=os.getenv("ANTHROPIC_API_KEY"),
|
||||
|
|
@ -17,10 +16,10 @@ anthropic = Anthropic(
|
|||
@ai_track("Tom Inference")
|
||||
@observe()
|
||||
async def get_tom_inference_conversational(
|
||||
chat_history: str, session_id: str, user_representation: str = "None"
|
||||
chat_history: str, user_representation: str = "None"
|
||||
) -> str:
|
||||
with sentry_sdk.start_transaction(op="tom-inference", name="ToM Inference"):
|
||||
messages = [
|
||||
messages: list[MessageParam] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
|
|
@ -77,21 +76,24 @@ async def get_tom_inference_conversational(
|
|||
temperature=0,
|
||||
messages=messages,
|
||||
)
|
||||
return message.content[0].text
|
||||
# skip blocks that are not text and return the first text block
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
return block.text
|
||||
raise RuntimeError("No text block returned by LLM")
|
||||
|
||||
|
||||
@ai_track("User Representation")
|
||||
@observe()
|
||||
async def get_user_representation_conversational(
|
||||
chat_history: str,
|
||||
session_id: str,
|
||||
user_representation: str = "None",
|
||||
tom_inference: str = "None",
|
||||
) -> str:
|
||||
with sentry_sdk.start_transaction(
|
||||
op="user-representation-inference", name="User Representation"
|
||||
):
|
||||
messages = [
|
||||
messages: list[MessageParam] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
|
|
@ -148,4 +150,8 @@ async def get_user_representation_conversational(
|
|||
temperature=0,
|
||||
messages=messages,
|
||||
)
|
||||
return message.content[0].text
|
||||
# skip blocks that are not text and return the first text block
|
||||
for block in message.content:
|
||||
if isinstance(block, TextBlock):
|
||||
return block.text
|
||||
raise RuntimeError("No text block returned by LLM")
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
import logging
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ... import crud, schemas
|
||||
from ...dependencies import tracked_db
|
||||
|
||||
|
|
@ -9,17 +7,17 @@ logger = logging.getLogger(__name__)
|
|||
|
||||
|
||||
class CollectionEmbeddingStore:
|
||||
def __init__(self, app_id: str, user_id: str, collection_id: str):
|
||||
self.app_id = app_id
|
||||
self.user_id = user_id
|
||||
self.collection_id = collection_id
|
||||
def __init__(self, workspace_name: str, peer_name: str, collection_name: str):
|
||||
self.workspace_name = workspace_name
|
||||
self.peer_name = peer_name
|
||||
self.collection_name = collection_name
|
||||
|
||||
async def save_facts(
|
||||
self,
|
||||
facts: list[str],
|
||||
replace_duplicates: bool = True,
|
||||
similarity_threshold: float = 0.85,
|
||||
message_id: str = None,
|
||||
message_id: int | None = None,
|
||||
) -> None:
|
||||
"""Save facts to the collection.
|
||||
|
||||
|
|
@ -33,14 +31,16 @@ class CollectionEmbeddingStore:
|
|||
# Create document with duplicate checking
|
||||
try:
|
||||
metadata = {}
|
||||
if message_id:
|
||||
if message_id is not None:
|
||||
metadata["message_id"] = message_id
|
||||
await crud.create_document(
|
||||
db,
|
||||
document=schemas.DocumentCreate(content=fact, metadata=metadata),
|
||||
app_id=self.app_id,
|
||||
user_id=self.user_id,
|
||||
collection_id=self.collection_id,
|
||||
document=schemas.DocumentCreate(
|
||||
content=fact, metadata=metadata
|
||||
),
|
||||
workspace_name=self.workspace_name,
|
||||
peer_name=self.peer_name,
|
||||
collection_name=self.collection_name,
|
||||
duplicate_threshold=1
|
||||
- similarity_threshold, # Convert similarity to distance
|
||||
)
|
||||
|
|
@ -64,9 +64,9 @@ class CollectionEmbeddingStore:
|
|||
async with tracked_db("embedding_store.get_relevant_facts") as db:
|
||||
documents = await crud.query_documents(
|
||||
db,
|
||||
app_id=self.app_id,
|
||||
user_id=self.user_id,
|
||||
collection_id=self.collection_id,
|
||||
workspace_name=self.workspace_name,
|
||||
peer_name=self.peer_name,
|
||||
collection_name=self.collection_name,
|
||||
query=query,
|
||||
max_distance=max_distance,
|
||||
top_k=top_k,
|
||||
|
|
@ -94,9 +94,9 @@ class CollectionEmbeddingStore:
|
|||
# Check for duplicates using the crud function
|
||||
duplicates = await crud.get_duplicate_documents(
|
||||
db,
|
||||
app_id=self.app_id,
|
||||
user_id=self.user_id,
|
||||
collection_id=self.collection_id,
|
||||
workspace_name=self.workspace_name,
|
||||
peer_name=self.peer_name,
|
||||
collection_name=self.collection_name,
|
||||
content=fact,
|
||||
similarity_threshold=similarity_threshold,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -26,7 +26,6 @@ MAX_FACT_DISTANCE = 0.85
|
|||
@observe()
|
||||
async def get_user_representation_long_term(
|
||||
chat_history: str,
|
||||
session_id: str,
|
||||
user_representation: str = "None",
|
||||
tom_inference: str = "None",
|
||||
facts: Optional[list[str]] = None,
|
||||
|
|
|
|||
|
|
@ -105,7 +105,6 @@ UPDATES:
|
|||
@observe()
|
||||
async def get_tom_inference_single_prompt(
|
||||
chat_history: str,
|
||||
session_id: str,
|
||||
user_representation: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
|
|
@ -153,7 +152,6 @@ async def get_tom_inference_single_prompt(
|
|||
@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,
|
||||
|
|
|
|||
18
src/main.py
18
src/main.py
|
|
@ -15,14 +15,11 @@ from sentry_sdk.integrations.starlette import StarletteIntegration
|
|||
from src.db import engine, request_context
|
||||
from src.exceptions import HonchoException
|
||||
from src.routers import (
|
||||
apps,
|
||||
collections,
|
||||
documents,
|
||||
keys,
|
||||
messages,
|
||||
metamessages,
|
||||
peers,
|
||||
sessions,
|
||||
users,
|
||||
workspaces,
|
||||
)
|
||||
from src.security import create_admin_jwt
|
||||
|
||||
|
|
@ -82,7 +79,7 @@ if SENTRY_ENABLED:
|
|||
# Sentry SDK's default behavior:
|
||||
# - Captures INFO+ level logs as breadcrumbs
|
||||
# - Captures ERROR+ level logs as Sentry events
|
||||
#
|
||||
#
|
||||
# For custom log levels, use the LoggingIntegration class:
|
||||
# sentry_sdk.init(..., integrations=[LoggingIntegration(level=logging.INFO, event_level=logging.ERROR)])
|
||||
sentry_sdk.init(
|
||||
|
|
@ -117,7 +114,7 @@ app = FastAPI(
|
|||
title="Honcho API",
|
||||
summary="The Identity Layer for the Agentic World",
|
||||
description="""Honcho is a platform for giving agents user-centric memory and social cognition""",
|
||||
version="1.1.0",
|
||||
version="2.0.0",
|
||||
contact={
|
||||
"name": "Plastic Labs",
|
||||
"url": "https://honcho.dev",
|
||||
|
|
@ -144,13 +141,10 @@ router = APIRouter(prefix="/apps/{app_id}/users/{user_id}")
|
|||
|
||||
add_pagination(app)
|
||||
|
||||
app.include_router(apps.router, prefix="/v1")
|
||||
app.include_router(users.router, prefix="/v1")
|
||||
app.include_router(workspaces.router, prefix="/v1")
|
||||
app.include_router(peers.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(collections.router, prefix="/v1")
|
||||
app.include_router(documents.router, prefix="/v1")
|
||||
app.include_router(keys.router, prefix="/v1")
|
||||
|
||||
|
||||
|
|
|
|||
383
src/models.py
383
src/models.py
|
|
@ -1,5 +1,7 @@
|
|||
import datetime
|
||||
from logging import getLogger
|
||||
|
||||
import tiktoken
|
||||
from dotenv import load_dotenv
|
||||
from nanoid import generate as generate_nanoid
|
||||
from pgvector.sqlalchemy import Vector
|
||||
|
|
@ -7,11 +9,16 @@ from sqlalchemy import (
|
|||
BigInteger,
|
||||
Boolean,
|
||||
CheckConstraint,
|
||||
Column,
|
||||
DateTime,
|
||||
ForeignKey,
|
||||
ForeignKeyConstraint,
|
||||
Identity,
|
||||
Index,
|
||||
Integer,
|
||||
Table,
|
||||
UniqueConstraint,
|
||||
event,
|
||||
text,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import JSONB, TEXT
|
||||
|
|
@ -22,265 +29,299 @@ from .db import Base
|
|||
|
||||
load_dotenv()
|
||||
|
||||
logger = getLogger(__name__)
|
||||
|
||||
class App(Base):
|
||||
__tablename__ = "apps"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
# Initialize tiktoken encoder for token counting for message content
|
||||
tokenizer = tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""Count tokens in a text string using tiktoken."""
|
||||
if not text:
|
||||
return 0
|
||||
try:
|
||||
return len(tokenizer.encode(text))
|
||||
except Exception as e:
|
||||
# Fallback: rough estimation (4 chars per token)
|
||||
logger.warning(
|
||||
f"Error counting tokens for text: {text[:50]}{'...' if len(text) > 50 else ''}, using fallback (4 chars per token). Error: {str(e)}"
|
||||
)
|
||||
return len(text) // 4
|
||||
|
||||
|
||||
# Association table for many-to-many relationship between sessions and peers
|
||||
session_peers_table = Table(
|
||||
"session_peers",
|
||||
Base.metadata,
|
||||
Column(
|
||||
"workspace_name",
|
||||
TEXT,
|
||||
ForeignKey("workspaces.name"),
|
||||
primary_key=True,
|
||||
nullable=False,
|
||||
),
|
||||
Column(
|
||||
"session_name",
|
||||
TEXT,
|
||||
primary_key=True,
|
||||
nullable=False,
|
||||
),
|
||||
Column("peer_name", TEXT, primary_key=True, nullable=False),
|
||||
Column("configuration", JSONB, default=dict),
|
||||
Column("internal_metadata", JSONB, default=dict),
|
||||
Column(
|
||||
"joined_at",
|
||||
DateTime(timezone=True),
|
||||
nullable=False,
|
||||
default=func.now(),
|
||||
),
|
||||
Column(
|
||||
"left_at",
|
||||
DateTime(timezone=True),
|
||||
nullable=True,
|
||||
),
|
||||
# Composite foreign key constraint for sessions
|
||||
ForeignKeyConstraint(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
),
|
||||
# Composite foreign key constraint for peers
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class Workspace(Base):
|
||||
__tablename__ = "workspaces"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True, unique=True)
|
||||
users = relationship("User", back_populates="app")
|
||||
peers = relationship("Peer", back_populates="workspace")
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
configuration: Mapped[dict] = mapped_column(JSONB, default=dict)
|
||||
|
||||
__table_args__ = (
|
||||
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("length(name) <= 512", name="name_length"),
|
||||
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
)
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
class Peer(Base):
|
||||
__tablename__ = "peers"
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
||||
app = relationship("App", back_populates="users")
|
||||
sessions = relationship("Session", back_populates="user")
|
||||
collections = relationship("Collection", back_populates="user")
|
||||
metamessages = relationship("Metamessage", back_populates="user")
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
)
|
||||
configuration: Mapped[dict] = mapped_column(JSONB, default=dict)
|
||||
|
||||
workspace = relationship("Workspace", back_populates="peers")
|
||||
sessions = relationship(
|
||||
"Session", secondary=session_peers_table, back_populates="peers"
|
||||
)
|
||||
collections = relationship("Collection", back_populates="peer")
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("name", "app_id", name="unique_name_app_user"),
|
||||
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
||||
UniqueConstraint("name", "workspace_name", name="unique_name_workspace_peer"),
|
||||
CheckConstraint("length(id) = 21", name="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"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
Index("idx_peers_workspace_lookup", "workspace_name", "name"),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"User(id={self.id}, app_id={self.app_id}, public_id={self.public_id} created_at={self.created_at}, h_metadata={self.h_metadata})"
|
||||
return f"Peer(id={self.id}, name={self.name}, workspace_name={self.workspace_name}, created_at={self.created_at}, h_metadata={self.h_metadata}, configuration={self.configuration})"
|
||||
|
||||
|
||||
class Session(Base):
|
||||
__tablename__ = "sessions"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, primary_key=True, default=generate_nanoid)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
is_active: Mapped[bool] = mapped_column(default=True)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
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)
|
||||
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
||||
user = relationship("User", back_populates="sessions")
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
)
|
||||
configuration: Mapped[dict] = mapped_column(JSONB, default=dict)
|
||||
|
||||
peers = relationship(
|
||||
"Peer", secondary=session_peers_table, 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"),
|
||||
UniqueConstraint("name", "workspace_name", name="unique_session_name"),
|
||||
CheckConstraint("length(name) <= 512", name="name_length"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Session(id={self.id}, user_id={self.user_id}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})"
|
||||
return f"Session(id={self.id}, name={self.name}, workspace_name={self.workspace_name}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})"
|
||||
|
||||
|
||||
class Message(Base):
|
||||
__tablename__ = "messages"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
session_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("sessions.public_id"), index=True
|
||||
)
|
||||
is_user: Mapped[bool]
|
||||
session_name: Mapped[str | None] = mapped_column(index=True, nullable=True)
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
token_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
|
||||
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
session = relationship("Session", back_populates="messages")
|
||||
metamessages = relationship("Metamessage", back_populates="message")
|
||||
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
||||
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
||||
peer_name: Mapped[str] = mapped_column(index=True)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
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"),
|
||||
# Composite foreign key constraint for sessions
|
||||
ForeignKeyConstraint(
|
||||
["session_name", "workspace_name"],
|
||||
["sessions.name", "sessions.workspace_name"],
|
||||
),
|
||||
# Composite foreign key constraint for peers
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
),
|
||||
Index(
|
||||
"idx_messages_session_lookup",
|
||||
"session_id",
|
||||
"session_name",
|
||||
"id",
|
||||
postgresql_include=["public_id", "is_user", "created_at"],
|
||||
postgresql_include=["id", "created_at"],
|
||||
),
|
||||
# Full text search index on content column
|
||||
Index(
|
||||
"idx_messages_content_gin",
|
||||
text("to_tsvector('english', content)"),
|
||||
postgresql_using="gin",
|
||||
),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Message(id={self.id}, session_id={self.session_id}, is_user={self.is_user}, content={self.content[10:]})"
|
||||
return f"Message(id={self.id}, session_name={self.session_name}, peer_name={self.peer_name}, content={self.content})"
|
||||
|
||||
|
||||
class Metamessage(Base):
|
||||
__tablename__ = "metamessages"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
label: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
|
||||
# Foreign keys - message_id is now optional
|
||||
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
||||
app_id: Mapped[str] = mapped_column(ForeignKey("apps.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()
|
||||
)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
|
||||
__table_args__ = (
|
||||
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"),
|
||||
CheckConstraint("length(label) <= 512", name="label_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",
|
||||
"label",
|
||||
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",
|
||||
"label",
|
||||
text("id DESC"),
|
||||
),
|
||||
Index(
|
||||
"idx_metamessages_session_lookup",
|
||||
"session_id",
|
||||
"label",
|
||||
text("id DESC"),
|
||||
),
|
||||
Index(
|
||||
"idx_metamessages_message_lookup",
|
||||
"message_id",
|
||||
"label",
|
||||
text("id DESC"),
|
||||
),
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Metamessages(id={self.id}, user_id={self.user_id}, session_id={self.session_id}, message_id={self.message_id}, label={self.label})"
|
||||
@event.listens_for(Message, "before_insert")
|
||||
def calculate_token_count_on_insert(_mapper, _connection, target):
|
||||
"""Calculate token count before inserting a new message."""
|
||||
target.token_count = count_tokens(target.content)
|
||||
|
||||
|
||||
class Collection(Base):
|
||||
__tablename__ = "collections"
|
||||
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default=dict)
|
||||
internal_metadata: Mapped[dict] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
documents = relationship(
|
||||
"Document", back_populates="collection", cascade="all, delete, delete-orphan"
|
||||
)
|
||||
user = relationship("User", back_populates="collections")
|
||||
user_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("users.public_id"), index=True
|
||||
peer = relationship("Peer", back_populates="collections")
|
||||
peer_name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
)
|
||||
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint("name", "user_id", name="unique_name_collection_user"),
|
||||
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
||||
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
||||
UniqueConstraint(
|
||||
"name", "peer_name", "workspace_name", name="unique_name_collection_peer"
|
||||
),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
CheckConstraint("length(name) <= 512", name="name_length"),
|
||||
# Composite foreign key constraint for peers
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class Document(Base):
|
||||
__tablename__ = "documents"
|
||||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
internal_metadata: Mapped[dict] = mapped_column(
|
||||
"internal_metadata", JSONB, default=dict
|
||||
)
|
||||
public_id: Mapped[str] = mapped_column(
|
||||
TEXT, index=True, unique=True, default=generate_nanoid
|
||||
)
|
||||
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
||||
content: Mapped[str] = mapped_column(TEXT)
|
||||
embedding = mapped_column(Vector(1536))
|
||||
created_at: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), index=True, default=func.now()
|
||||
)
|
||||
|
||||
collection_id: Mapped[str] = mapped_column(
|
||||
TEXT, ForeignKey("collections.public_id"), index=True
|
||||
collection_name: Mapped[str] = mapped_column(TEXT, index=True)
|
||||
peer_name: Mapped[str] = mapped_column(index=True)
|
||||
workspace_name: Mapped[str] = mapped_column(
|
||||
ForeignKey("workspaces.name"), index=True
|
||||
)
|
||||
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
||||
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
||||
collection = relationship("Collection", back_populates="documents")
|
||||
|
||||
__table_args__ = (
|
||||
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
||||
CheckConstraint("length(id) = 21", name="id_length"),
|
||||
CheckConstraint("length(content) <= 65535", name="content_length"),
|
||||
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
||||
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
|
||||
# Composite foreign key constraint for collections
|
||||
ForeignKeyConstraint(
|
||||
["collection_name", "peer_name", "workspace_name"],
|
||||
["collections.name", "collections.peer_name", "collections.workspace_name"],
|
||||
),
|
||||
# Composite foreign key constraint for peers
|
||||
ForeignKeyConstraint(
|
||||
["peer_name", "workspace_name"],
|
||||
["peers.name", "peers.workspace_name"],
|
||||
),
|
||||
# HNSW index on embedding column
|
||||
Index(
|
||||
"idx_documents_embedding_hnsw",
|
||||
"embedding",
|
||||
postgresql_using="hnsw", # HNSW index type
|
||||
postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters
|
||||
postgresql_ops={"embedding": "vector_cosine_ops"}, # Cosine distance operator
|
||||
postgresql_using="hnsw", # HNSW index type
|
||||
postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters
|
||||
postgresql_ops={
|
||||
"embedding": "vector_cosine_ops"
|
||||
}, # Cosine distance operator
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -290,7 +331,9 @@ class QueueItem(Base):
|
|||
id: Mapped[int] = mapped_column(
|
||||
BigInteger, Identity(), primary_key=True, autoincrement=True
|
||||
)
|
||||
session_id: Mapped[int] = mapped_column(ForeignKey("sessions.id"), index=True)
|
||||
session_id: Mapped[str] = mapped_column(
|
||||
ForeignKey("sessions.id"), index=True, nullable=True
|
||||
)
|
||||
payload: Mapped[dict] = mapped_column(JSONB, nullable=False)
|
||||
processed: Mapped[bool] = mapped_column(Boolean, default=False)
|
||||
|
||||
|
|
@ -298,7 +341,27 @@ class QueueItem(Base):
|
|||
class ActiveQueueSession(Base):
|
||||
__tablename__ = "active_queue_sessions"
|
||||
|
||||
session_id: Mapped[int] = mapped_column(BigInteger, primary_key=True, index=True)
|
||||
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
|
||||
session_id: Mapped[str | None] = mapped_column(
|
||||
ForeignKey("sessions.id"), nullable=True
|
||||
)
|
||||
sender_name: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
target_name: Mapped[str | None] = mapped_column(TEXT, nullable=True)
|
||||
task_type: Mapped[str] = mapped_column(TEXT)
|
||||
last_updated: Mapped[datetime.datetime] = mapped_column(
|
||||
DateTime(timezone=True), default=func.now(), onupdate=func.now()
|
||||
)
|
||||
|
||||
__table_args__ = (
|
||||
UniqueConstraint(
|
||||
"session_id",
|
||||
"sender_name",
|
||||
"target_name",
|
||||
"task_type",
|
||||
name="unique_active_queue_session",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class SessionPeer(Base):
|
||||
__table__ = session_peers_table
|
||||
|
|
|
|||
|
|
@ -1,138 +0,0 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
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 JWTParams, require_auth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps",
|
||||
tags=["apps"],
|
||||
)
|
||||
|
||||
|
||||
@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.post(
|
||||
"/list",
|
||||
response_model=Page[schemas.App],
|
||||
dependencies=[Depends(require_auth(admin=True))],
|
||||
)
|
||||
async def get_all_apps(
|
||||
options: Optional[schemas.AppGet] = Body(
|
||||
None, 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"""
|
||||
filter_param = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}:
|
||||
filter_param = None
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_all_apps(
|
||||
reverse=reverse,
|
||||
filter=filter_param,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@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, 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,
|
||||
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)
|
||||
return app
|
||||
except ResourceNotFoundException:
|
||||
# App doesn't exist, create it
|
||||
app = await create_app(db=db, app=schemas.AppCreate(name=name))
|
||||
return app
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{app_id}",
|
||||
response_model=schemas.App,
|
||||
dependencies=[Depends(require_auth(app_id="app_id"))],
|
||||
)
|
||||
async def update_app(
|
||||
app_id: str = Path(..., description="ID of the app to update"),
|
||||
app: schemas.AppUpdate = Body(..., description="Updated app parameters"),
|
||||
db=db,
|
||||
):
|
||||
"""Update an App"""
|
||||
# ResourceNotFoundException will be caught by global handler if app not found
|
||||
honcho_app = await crud.update_app(db, app_id=app_id, app=app)
|
||||
return honcho_app
|
||||
|
|
@ -1,190 +0,0 @@
|
|||
from typing import Optional
|
||||
|
||||
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
|
||||
from src.security import JWTParams, require_auth
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps/{app_id}/users/{user_id}/collections",
|
||||
tags=["collections"],
|
||||
)
|
||||
|
||||
|
||||
@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 = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
options: Optional[schemas.CollectionGet] = Body(
|
||||
None, 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"""
|
||||
filter_param = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}: # Explicitly check for empty dict
|
||||
filter_param = None
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_collections(
|
||||
app_id=app_id, user_id=user_id, filter=filter_param, reverse=reverse
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@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 = 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"""
|
||||
honcho_collection = await crud.get_collection_by_name(
|
||||
db, app_id=app_id, user_id=user_id, name=name
|
||||
)
|
||||
return honcho_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 = 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"""
|
||||
# ValidationException will be caught by global handler if collection is invalid
|
||||
# ConflictException will be caught by global handler if collection name already exists
|
||||
return await crud.create_collection(
|
||||
db, collection=collection, app_id=app_id, user_id=user_id
|
||||
)
|
||||
|
||||
|
||||
@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 = 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"
|
||||
# ResourceNotFoundException will be caught by global handler if collection not found
|
||||
# ValidationException will be caught by global handler if update data is invalid
|
||||
honcho_collection = await crud.update_collection(
|
||||
db,
|
||||
collection=collection,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
)
|
||||
return honcho_collection
|
||||
|
||||
|
||||
@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 = 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"""
|
||||
# ResourceNotFoundException will be caught by global handler if collection not found
|
||||
await crud.delete_collection(
|
||||
db, app_id=app_id, user_id=user_id, collection_id=collection_id
|
||||
)
|
||||
return {"message": "Collection deleted successfully"}
|
||||
|
|
@ -1,212 +0,0 @@
|
|||
import logging
|
||||
from collections.abc import Sequence
|
||||
from typing import Optional
|
||||
|
||||
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 require_auth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps/{app_id}/users/{user_id}/collections/{collection_id}/documents",
|
||||
tags=["documents"],
|
||||
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 = 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: Optional[schemas.DocumentGet] = Body(
|
||||
None, 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"""
|
||||
filter_param = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}: # Explicitly check for empty dict
|
||||
filter_param = None
|
||||
|
||||
try:
|
||||
documents_query = await crud.get_documents(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
filter=filter_param,
|
||||
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
|
||||
|
||||
|
||||
@router.get("/{document_id}", response_model=schemas.Document)
|
||||
async def get_document(
|
||||
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"""
|
||||
honcho_document = await crud.get_document(
|
||||
db,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
document_id=document_id,
|
||||
)
|
||||
return honcho_document
|
||||
|
||||
|
||||
@router.post("/query", response_model=Sequence[schemas.Document])
|
||||
async def query_documents(
|
||||
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"""
|
||||
|
||||
try:
|
||||
top_k = options.top_k
|
||||
filter = options.filter
|
||||
if options.filter == {}:
|
||||
filter = None
|
||||
|
||||
documents = await crud.query_documents(
|
||||
db=db,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
query=options.query,
|
||||
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)}"
|
||||
)
|
||||
raise ValidationException("Error querying documents") from e
|
||||
|
||||
|
||||
@router.post("", response_model=schemas.Document)
|
||||
async def create_document(
|
||||
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"""
|
||||
try:
|
||||
document_obj = await crud.create_document(
|
||||
db,
|
||||
document=document,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{document_id}",
|
||||
response_model=schemas.Document,
|
||||
)
|
||||
async def update_document(
|
||||
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}"
|
||||
)
|
||||
raise ValidationException("Content and metadata cannot both be None")
|
||||
|
||||
try:
|
||||
updated_document = await crud.update_document(
|
||||
db,
|
||||
document=document,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
document_id=document_id,
|
||||
)
|
||||
logger.info(f"Document {document_id} updated successfully")
|
||||
return updated_document
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to update document {document_id}: {str(e)}")
|
||||
raise ResourceNotFoundException("Collection or document not found") from e
|
||||
|
||||
|
||||
@router.delete("/{document_id}")
|
||||
async def delete_document(
|
||||
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"""
|
||||
response = await crud.delete_document(
|
||||
db,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
collection_id=collection_id,
|
||||
document_id=document_id,
|
||||
)
|
||||
if response:
|
||||
logger.info(f"Document {document_id} deleted successfully")
|
||||
return {"message": "Document deleted successfully"}
|
||||
else:
|
||||
logger.warning(f"Document {document_id} not found or could not be deleted")
|
||||
raise ResourceNotFoundException("Document not found or does not belong to user")
|
||||
|
|
@ -24,14 +24,13 @@ router = APIRouter(
|
|||
|
||||
@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"),
|
||||
workspace_id: str | None = Query(
|
||||
None, description="ID of the workspace to scope the key to"
|
||||
),
|
||||
peer_id: str | None = Query(None, description="ID of the peer 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"""
|
||||
|
|
@ -39,18 +38,17 @@ async def create_key(
|
|||
raise DisabledException()
|
||||
|
||||
# Validate that at least one parameter is provided for proper scoping
|
||||
if not any([app_id, user_id, session_id, collection_id]):
|
||||
if not any([workspace_id, peer_id, session_id]):
|
||||
raise ValidationException(
|
||||
"At least one of app_id, user_id, session_id, or collection_id must be provided"
|
||||
"At least one of workspace_id, peer_id, or session_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,
|
||||
w=workspace_id,
|
||||
p=peer_id,
|
||||
s=session_id,
|
||||
)
|
||||
)
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
from typing import Any, Optional
|
||||
|
||||
from fastapi import APIRouter, BackgroundTasks, Body, Depends, Path, Query
|
||||
from fastapi_pagination import Page
|
||||
|
|
@ -16,17 +16,44 @@ from src.security import require_auth
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps/{app_id}/users/{user_id}/sessions/{session_id}/messages",
|
||||
prefix="/workspaces/{workspace_id}/sessions/{session_id}/messages",
|
||||
tags=["messages"],
|
||||
dependencies=[
|
||||
Depends(
|
||||
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
|
||||
)
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
async def enqueue(payload: dict | list[dict]):
|
||||
def create_processed_payload(
|
||||
message: dict,
|
||||
sender_name: Optional[str],
|
||||
target_name: Optional[str],
|
||||
task_type: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Create a processed payload from a message for queue processing.
|
||||
|
||||
Args:
|
||||
message: The original message dictionary
|
||||
sender_name: Name of the message sender
|
||||
target_name: Name of the target peer
|
||||
task_type: Type of task ('representation' or 'summary')
|
||||
|
||||
Returns:
|
||||
Processed payload dictionary ready for queue processing
|
||||
"""
|
||||
processed_payload = {
|
||||
k: str(v) if isinstance(v, str) else v for k, v in message.items()
|
||||
}
|
||||
# Remove peer_name from payload
|
||||
processed_payload.pop("peer_name", None) # Use None as default to avoid KeyError
|
||||
processed_payload["sender_name"] = sender_name
|
||||
processed_payload["target_name"] = target_name
|
||||
processed_payload["task_type"] = task_type
|
||||
return processed_payload
|
||||
|
||||
|
||||
async def enqueue(payload: list[dict]):
|
||||
"""
|
||||
Add message(s) to the deriver queue for processing.
|
||||
|
||||
|
|
@ -37,189 +64,226 @@ async def enqueue(payload: dict | list[dict]):
|
|||
# Use the get_db dependency to ensure proper transaction handling
|
||||
async with tracked_db("message_enqueue") as db_session:
|
||||
try:
|
||||
if isinstance(payload, list):
|
||||
if not payload: # Empty list check
|
||||
logger.debug("Empty payload list, skipping enqueue")
|
||||
return
|
||||
# Determine if batch or single processing
|
||||
if not payload: # Empty list check
|
||||
logger.debug("Empty payload list, skipping enqueue")
|
||||
return
|
||||
logger.debug(f"Enqueueing batch of {len(payload)} messages")
|
||||
workspace_name = payload[0]["workspace_name"]
|
||||
session_name = payload[0]["session_name"]
|
||||
|
||||
logger.debug(f"Enqueueing batch of {len(payload)} messages")
|
||||
# Case 1: session_name is None — only create representation for peer
|
||||
if session_name is None:
|
||||
peer_name = payload[0]["peer_name"]
|
||||
logger.info(
|
||||
"Session name is None, creating single representation queue items"
|
||||
)
|
||||
peer = await crud.get_or_create_peers(
|
||||
db_session,
|
||||
workspace_name=workspace_name,
|
||||
peers=[schemas.PeerCreate(name=peer_name)],
|
||||
)
|
||||
peer = peer[0]
|
||||
|
||||
# Check session once since all messages are for same session
|
||||
try:
|
||||
session = await crud.get_session(
|
||||
db_session,
|
||||
app_id=payload[0]["app_id"],
|
||||
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"
|
||||
)
|
||||
return
|
||||
|
||||
logger.debug(f"Session {session.public_id} found for batch enqueue")
|
||||
|
||||
# Check if deriver is disabled for this session
|
||||
if (
|
||||
session.h_metadata.get("deriver_disabled") is not None
|
||||
and session.h_metadata.get("deriver_disabled") is not False
|
||||
):
|
||||
# Cast configuration to PeerConfig and check observe_me
|
||||
peer_config = (
|
||||
schemas.PeerConfig(**peer.configuration)
|
||||
if peer.configuration
|
||||
else schemas.PeerConfig()
|
||||
)
|
||||
if not peer_config.observe_me:
|
||||
logger.info(
|
||||
f"Deriver is disabled for session {session.public_id}, skipping enqueue"
|
||||
f"Peer {peer_name} has observe_me=False, skipping enqueue"
|
||||
)
|
||||
return
|
||||
|
||||
# Process all payloads
|
||||
queue_records = [
|
||||
queue_records: list[dict[str, Any]] = []
|
||||
|
||||
for message in payload:
|
||||
processed_payload = create_processed_payload(
|
||||
message=message,
|
||||
sender_name=message["peer_name"],
|
||||
target_name=message["peer_name"],
|
||||
task_type="representation",
|
||||
)
|
||||
queue_records.append(
|
||||
{
|
||||
"payload": processed_payload,
|
||||
"session_id": None,
|
||||
}
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"Inserting {len(queue_records)} queue records for None session"
|
||||
)
|
||||
stmt = insert(QueueItem).returning(QueueItem)
|
||||
await db_session.execute(stmt, queue_records)
|
||||
await db_session.commit()
|
||||
logger.info(
|
||||
f"Successfully enqueued {len(queue_records)} messages with None session"
|
||||
)
|
||||
return
|
||||
|
||||
# Case 2: Normal session processing
|
||||
session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
session=schemas.SessionCreate(name=session_name),
|
||||
workspace_name=workspace_name,
|
||||
)
|
||||
|
||||
# Check if deriver is disabled for this session
|
||||
deriver_disabled = (
|
||||
session.configuration.get("deriver_disabled") is not None
|
||||
and session.configuration.get("deriver_disabled") is not False
|
||||
)
|
||||
|
||||
configuration_query = await crud.get_session_peer_configuration(
|
||||
workspace_name=workspace_name, session_name=session_name
|
||||
)
|
||||
peers_with_configuration_result = await db_session.execute(
|
||||
configuration_query
|
||||
)
|
||||
peers_with_configuration_list = peers_with_configuration_result.all()
|
||||
peers_with_configuration = {
|
||||
row.peer_name: [row.peer_configuration, row.session_peer_configuration]
|
||||
for row in peers_with_configuration_list
|
||||
}
|
||||
|
||||
# Process all payloads - create multiple queue items per message
|
||||
queue_records = []
|
||||
|
||||
for message in payload:
|
||||
if deriver_disabled:
|
||||
# still create a summary queue item for the session
|
||||
processed_payload = create_processed_payload(
|
||||
message=message,
|
||||
sender_name=None,
|
||||
target_name=None,
|
||||
task_type="summary",
|
||||
)
|
||||
queue_records.append(
|
||||
{
|
||||
"payload": processed_payload,
|
||||
"session_id": session.id,
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
sender_name = message["peer_name"]
|
||||
|
||||
sender_session_peer_config = (
|
||||
schemas.SessionPeerConfig(
|
||||
**peers_with_configuration[sender_name][1]
|
||||
)
|
||||
if peers_with_configuration[sender_name][1]
|
||||
else None
|
||||
)
|
||||
sender_peer_config = (
|
||||
schemas.PeerConfig(**peers_with_configuration[sender_name][0])
|
||||
if peers_with_configuration[sender_name][0]
|
||||
else schemas.PeerConfig()
|
||||
)
|
||||
|
||||
observe_me = (
|
||||
sender_session_peer_config.observe_me
|
||||
if sender_session_peer_config
|
||||
else sender_peer_config.observe_me
|
||||
)
|
||||
if not observe_me:
|
||||
continue
|
||||
|
||||
# Handle working representation for sender
|
||||
processed_payload = create_processed_payload(
|
||||
message=message,
|
||||
sender_name=sender_name,
|
||||
target_name=sender_name,
|
||||
task_type="representation",
|
||||
)
|
||||
|
||||
queue_records.append(
|
||||
{
|
||||
"payload": {
|
||||
k: str(v) if isinstance(v, str) else v for k, v in p.items()
|
||||
},
|
||||
"payload": processed_payload,
|
||||
"session_id": session.id,
|
||||
}
|
||||
for p in payload
|
||||
]
|
||||
)
|
||||
for peer_name, configuration in peers_with_configuration.items():
|
||||
session_peer_config = (
|
||||
schemas.SessionPeerConfig(**configuration[1])
|
||||
if configuration[1]
|
||||
else None
|
||||
)
|
||||
|
||||
logger.debug(f"Inserting {len(queue_records)} queue records")
|
||||
if peer_name != sender_name:
|
||||
# Handle local representation for other peers
|
||||
if (
|
||||
session_peer_config is None
|
||||
or not session_peer_config.observe_others
|
||||
):
|
||||
continue
|
||||
else:
|
||||
# Create local representation for peer
|
||||
processed_payload = create_processed_payload(
|
||||
message=message,
|
||||
sender_name=sender_name,
|
||||
target_name=peer_name,
|
||||
task_type="representation",
|
||||
)
|
||||
|
||||
queue_records.append(
|
||||
{
|
||||
"payload": processed_payload,
|
||||
"session_id": session.id,
|
||||
}
|
||||
)
|
||||
|
||||
logger.debug(f"Inserting {len(queue_records)} queue records")
|
||||
|
||||
if len(queue_records) > 0:
|
||||
# Use insert to maintain order
|
||||
stmt = insert(QueueItem).returning(QueueItem)
|
||||
await db_session.execute(stmt, queue_records)
|
||||
await db_session.commit()
|
||||
logger.info(f"Successfully enqueued batch of {len(payload)} messages")
|
||||
return
|
||||
else:
|
||||
# Single message enqueue
|
||||
logger.debug(
|
||||
f"Enqueueing single message for session {payload['session_id']}"
|
||||
)
|
||||
|
||||
try:
|
||||
session = await crud.get_session(
|
||||
db_session,
|
||||
app_id=payload["app_id"],
|
||||
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"
|
||||
)
|
||||
return
|
||||
|
||||
# Check if deriver is disabled for this session
|
||||
deriver_disabled = session.h_metadata.get("deriver_disabled")
|
||||
if deriver_disabled is not None and deriver_disabled is not False:
|
||||
logger.info(
|
||||
f"Deriver is disabled for session {payload['session_id']}, skipping enqueue"
|
||||
)
|
||||
return
|
||||
|
||||
processed_payload = {
|
||||
k: str(v) if isinstance(v, str) else v for k, v in payload.items()
|
||||
}
|
||||
# Use insert for consistency
|
||||
stmt = (
|
||||
insert(QueueItem)
|
||||
.values(payload=processed_payload, session_id=session.id)
|
||||
.returning(QueueItem)
|
||||
)
|
||||
await db_session.execute(stmt)
|
||||
await db_session.commit()
|
||||
logger.info(
|
||||
f"Successfully enqueued message for session {payload['session_id']}"
|
||||
f"Successfully enqueued {len(payload)} messages with {len(queue_records)} total queue items"
|
||||
)
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to enqueue message: {str(e)}", exc_info=True)
|
||||
logger.error(f"Failed to enqueue messages: {str(e)}", exc_info=True)
|
||||
if os.getenv("SENTRY_ENABLED", "False").lower() == "true":
|
||||
import sentry_sdk
|
||||
|
||||
sentry_sdk.capture_exception(e)
|
||||
|
||||
|
||||
@router.post("", response_model=schemas.Message)
|
||||
async def create_message_for_session(
|
||||
@router.post("/", response_model=list[schemas.Message])
|
||||
async def create_messages_for_session(
|
||||
background_tasks: BackgroundTasks,
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
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"""
|
||||
try:
|
||||
honcho_message = await crud.create_message(
|
||||
db, message=message, app_id=app_id, user_id=user_id, session_id=session_id
|
||||
)
|
||||
|
||||
# Prepare message payload for background processing
|
||||
payload = {
|
||||
"app_id": app_id,
|
||||
"user_id": user_id,
|
||||
"session_id": session_id,
|
||||
"message_id": honcho_message.public_id,
|
||||
"is_user": honcho_message.is_user,
|
||||
"content": honcho_message.content,
|
||||
"metadata": honcho_message.h_metadata,
|
||||
}
|
||||
|
||||
# Queue message for background processing
|
||||
background_tasks.add_task(enqueue, payload) # type: ignore
|
||||
logger.info(
|
||||
f"Message {honcho_message.public_id} created and queued for processing"
|
||||
)
|
||||
|
||||
return honcho_message
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to create message for session {session_id}: {str(e)}")
|
||||
raise ResourceNotFoundException("Session not found") from e
|
||||
|
||||
|
||||
@router.post("/batch", response_model=list[schemas.Message])
|
||||
async def create_batch_messages_for_session(
|
||||
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(
|
||||
messages: schemas.MessageBatchCreate = Body(
|
||||
..., description="Batch of messages to create"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
workspace_name, session_name = workspace_id, session_id
|
||||
"""Bulk create messages for a session while maintaining order. Maximum 100 messages per batch."""
|
||||
try:
|
||||
created_messages = await crud.create_messages(
|
||||
db,
|
||||
messages=batch.messages,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
messages=messages.messages,
|
||||
workspace_name=workspace_name,
|
||||
session_name=session_name,
|
||||
)
|
||||
|
||||
# Create payloads for all messages
|
||||
payloads = [
|
||||
{
|
||||
"app_id": app_id,
|
||||
"user_id": user_id,
|
||||
"session_id": session_id,
|
||||
"message_id": message.public_id,
|
||||
"is_user": message.is_user,
|
||||
"workspace_name": workspace_name,
|
||||
"session_name": session_name,
|
||||
"message_id": message.id,
|
||||
"content": message.content,
|
||||
"metadata": message.h_metadata,
|
||||
"peer_name": message.peer_name,
|
||||
}
|
||||
for message in created_messages
|
||||
]
|
||||
|
|
@ -240,8 +304,7 @@ async def create_batch_messages_for_session(
|
|||
|
||||
@router.post("/list", response_model=Page[schemas.Message])
|
||||
async def get_messages(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
options: Optional[schemas.MessageGet] = Body(
|
||||
None, description="Filtering options for the messages list"
|
||||
|
|
@ -260,9 +323,8 @@ async def get_messages(
|
|||
filter = None
|
||||
|
||||
messages_query = await crud.get_messages(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
filter=filter,
|
||||
reverse=reverse,
|
||||
)
|
||||
|
|
@ -275,15 +337,14 @@ async def get_messages(
|
|||
|
||||
@router.get("/{message_id}", response_model=schemas.Message)
|
||||
async def get_message(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
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"""
|
||||
honcho_message = await crud.get_message(
|
||||
db, app_id=app_id, session_id=session_id, user_id=user_id, message_id=message_id
|
||||
db, workspace_name=workspace_id, session_name=session_id, message_id=message_id
|
||||
)
|
||||
if honcho_message is None:
|
||||
logger.warning(f"Message {message_id} not found in session {session_id}")
|
||||
|
|
@ -293,8 +354,7 @@ async def get_message(
|
|||
|
||||
@router.put("/{message_id}", response_model=schemas.Message)
|
||||
async def update_message(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
message_id: str = Path(..., description="ID of the message to update"),
|
||||
message: schemas.MessageUpdate = Body(
|
||||
|
|
@ -307,9 +367,8 @@ async def update_message(
|
|||
updated_message = await crud.update_message(
|
||||
db,
|
||||
message=message,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
message_id=message_id,
|
||||
)
|
||||
logger.info(f"Message {message_id} updated successfully")
|
||||
|
|
|
|||
|
|
@ -1,154 +0,0 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
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 require_auth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps/{app_id}/users/{user_id}/metamessages",
|
||||
tags=["metamessages"],
|
||||
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 = 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,
|
||||
):
|
||||
"""
|
||||
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,
|
||||
)
|
||||
logger.info(f"Metamessage created successfully for user {user_id}")
|
||||
return metamessage_obj
|
||||
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 = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
options: Optional[schemas.MetamessageGet] = Body(
|
||||
None, description="Filtering options for the metamessages list"
|
||||
),
|
||||
reverse: Optional[bool] = Query(
|
||||
False, description="Whether to reverse the order of results"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""
|
||||
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 label
|
||||
- Filter by metadata: Provide filter object
|
||||
"""
|
||||
session_id_param = None
|
||||
message_id_param = None
|
||||
label_param = None
|
||||
filter_param = None
|
||||
|
||||
if options:
|
||||
if hasattr(options, "session_id") and options.session_id:
|
||||
session_id_param = options.session_id
|
||||
if hasattr(options, "message_id") and options.message_id:
|
||||
message_id_param = options.message_id
|
||||
if hasattr(options, "label") and options.label:
|
||||
label_param = options.label
|
||||
if hasattr(options, "filter") and options.filter:
|
||||
filter_param = options.filter
|
||||
if filter_param == {}: # Explicitly check for empty dict
|
||||
filter_param = None
|
||||
|
||||
try:
|
||||
metamessages_query = await crud.get_metamessages(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id_param,
|
||||
message_id=message_id_param,
|
||||
label=label_param,
|
||||
filter=filter_param,
|
||||
reverse=reverse,
|
||||
)
|
||||
return await paginate(db, metamessages_query)
|
||||
except (ResourceNotFoundException, ValidationException) as e:
|
||||
logger.warning(f"Failed to get metamessages: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{metamessage_id}",
|
||||
response_model=schemas.Metamessage,
|
||||
)
|
||||
async def get_metamessage(
|
||||
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,
|
||||
user_id=user_id,
|
||||
metamessage_id=metamessage_id,
|
||||
)
|
||||
if honcho_metamessage is None:
|
||||
logger.warning(f"Metamessage {metamessage_id} not found")
|
||||
raise ResourceNotFoundException(
|
||||
f"Metamessage with ID {metamessage_id} not found"
|
||||
)
|
||||
return honcho_metamessage
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{metamessage_id}",
|
||||
response_model=schemas.Metamessage,
|
||||
)
|
||||
async def update_metamessage(
|
||||
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 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,
|
||||
metamessage_id=metamessage_id,
|
||||
)
|
||||
logger.info(f"Metamessage {metamessage_id} updated successfully")
|
||||
return updated_metamessage
|
||||
except (ResourceNotFoundException, ValidationException) as e:
|
||||
logger.warning(f"Failed to update metamessage {metamessage_id}: {str(e)}")
|
||||
raise
|
||||
|
|
@ -0,0 +1,315 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import (
|
||||
APIRouter,
|
||||
BackgroundTasks,
|
||||
Body,
|
||||
Depends,
|
||||
Path,
|
||||
Query,
|
||||
)
|
||||
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 (
|
||||
AuthenticationException,
|
||||
ResourceNotFoundException,
|
||||
)
|
||||
from src.routers.messages import enqueue
|
||||
from src.security import JWTParams, require_auth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/workspaces/{workspace_id}/peers",
|
||||
tags=["peers"],
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/list",
|
||||
response_model=Page[schemas.Peer],
|
||||
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
|
||||
)
|
||||
async def get_peers(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
options: Optional[schemas.PeerGet] = Body(
|
||||
None, description="Filtering options for the peers list"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Get All Peers for a Workspace"""
|
||||
filter_param = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}:
|
||||
filter_param = None
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_peers(workspace_name=workspace_id, filter=filter_param),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"",
|
||||
response_model=schemas.Peer,
|
||||
)
|
||||
async def get_or_create_peer(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer: schemas.PeerCreate = Body(..., description="Peer creation parameters"),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
db=db,
|
||||
):
|
||||
"""
|
||||
Get a Peer by ID
|
||||
|
||||
If peer_id is provided as a query parameter, it uses that (must match JWT workspace_id).
|
||||
Otherwise, it uses the peer_id from the JWT token.
|
||||
"""
|
||||
# validate workspace query param
|
||||
if not jwt_params.ad and jwt_params.w is not None and jwt_params.w != workspace_id:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
|
||||
if peer.name:
|
||||
if not jwt_params.ad and jwt_params.p is not None and jwt_params.p != peer.name:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
else:
|
||||
# Use peer_id from JWT
|
||||
if not jwt_params.p:
|
||||
raise AuthenticationException("Peer ID not found in query parameter or JWT")
|
||||
peer.name = jwt_params.p
|
||||
peer = (
|
||||
await crud.get_or_create_peers(db, workspace_name=workspace_id, peers=[peer])
|
||||
)[0]
|
||||
return peer
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{peer_id}",
|
||||
response_model=schemas.Peer,
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def update_peer(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer to update"),
|
||||
peer: schemas.PeerUpdate = Body(..., description="Updated peer parameters"),
|
||||
db=db,
|
||||
):
|
||||
"""Update a Peer's name and/or metadata"""
|
||||
updated_peer = await crud.update_peer(
|
||||
db, workspace_name=workspace_id, peer_name=peer_id, peer=peer
|
||||
)
|
||||
return updated_peer
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{peer_id}/sessions",
|
||||
response_model=Page[schemas.Session],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def get_sessions_for_peer(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
options: Optional[schemas.SessionGet] = Body(
|
||||
None, description="Filtering options for the sessions list"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Get All Sessions for a Peer"""
|
||||
filter_param = None
|
||||
is_active = True
|
||||
|
||||
if options:
|
||||
if hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}:
|
||||
filter_param = None
|
||||
if hasattr(options, "is_active"):
|
||||
is_active = options.is_active
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_sessions_for_peer(
|
||||
workspace_name=workspace_id,
|
||||
peer_name=peer_id,
|
||||
is_active=is_active,
|
||||
filter=filter_param,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{peer_id}/chat",
|
||||
response_model=schemas.DialecticResponse,
|
||||
responses={
|
||||
200: {
|
||||
"description": "Response to a question informed by Honcho's User Representation",
|
||||
"content": {"text/event-stream": {}},
|
||||
},
|
||||
},
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def chat(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
options: schemas.DialecticOptions = Body(
|
||||
..., description="Dialectic Endpoint Parameters"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
# Get or create the peer to ensure it exists
|
||||
await crud.get_or_create_peers(
|
||||
db, workspace_name=workspace_id, peers=[schemas.PeerCreate(name=peer_id)]
|
||||
)
|
||||
|
||||
response = await agent.chat(
|
||||
workspace_id, peer_id, options.session_id, options.queries, options.stream
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{peer_id}/messages",
|
||||
response_model=list[schemas.Message],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def create_messages_for_peer(
|
||||
background_tasks: BackgroundTasks,
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
messages: schemas.MessageBatchCreate = Body(
|
||||
..., description="Batch of messages to create"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Create messages for a peer"""
|
||||
workspace_name, peer_name = workspace_id, peer_id
|
||||
"""Bulk create messages for a peer while maintaining order."""
|
||||
try:
|
||||
created_messages = await crud.create_messages_for_peer(
|
||||
db,
|
||||
messages=messages.messages,
|
||||
workspace_name=workspace_name,
|
||||
peer_name=peer_name,
|
||||
)
|
||||
|
||||
# Create payloads for all messages
|
||||
payloads = [
|
||||
{
|
||||
"workspace_name": workspace_name,
|
||||
"session_name": None,
|
||||
"message_id": message.id,
|
||||
"content": message.content,
|
||||
"peer_name": message.peer_name,
|
||||
}
|
||||
for message in created_messages
|
||||
]
|
||||
|
||||
# Enqueue all messages in one call
|
||||
background_tasks.add_task(enqueue, payloads) # type: ignore
|
||||
logger.info(
|
||||
f"Batch of {len(created_messages)} messages created and queued for processing"
|
||||
)
|
||||
|
||||
return created_messages
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to create batch messages for peer {peer_id}: {str(e)}")
|
||||
raise ResourceNotFoundException("Peer not found") from e
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{peer_id}/messages/list",
|
||||
response_model=Page[schemas.Message],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def get_messages_for_peer(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
options: Optional[schemas.MessageGet] = Body(
|
||||
None, 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 peer"""
|
||||
try:
|
||||
filter = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter = options.filter
|
||||
if filter == {}:
|
||||
filter = None
|
||||
|
||||
messages_query = await crud.get_messages_for_peer(
|
||||
workspace_name=workspace_id,
|
||||
peer_name=peer_id,
|
||||
filter=filter,
|
||||
reverse=reverse,
|
||||
)
|
||||
|
||||
return await paginate(db, messages_query)
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to get messages for peer {peer_id}: {str(e)}")
|
||||
raise ResourceNotFoundException("Peer not found") from e
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{peer_id}/representation",
|
||||
response_model=dict[str, object],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def get_working_representation(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
options: schemas.PeerRepresentationGet = Body(
|
||||
..., description="Options for getting the peer representation"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Get a peer's working representation for a session.
|
||||
|
||||
If a session_id is provided in the body, we get the working representation of the peer in that session.
|
||||
|
||||
In the current implementation, we don't offer representations of `target` so that parameter is ignored.
|
||||
Future releases will allow for this.
|
||||
"""
|
||||
representation = await crud.get_working_representation(
|
||||
db, workspace_id, peer_id, options.session_id
|
||||
)
|
||||
return {"representation": representation}
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{peer_id}/search",
|
||||
response_model=Page[schemas.Message],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", peer_name="peer_id"))
|
||||
],
|
||||
)
|
||||
async def search_peer(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
query: str = Body(..., description="Search query"),
|
||||
db=db,
|
||||
):
|
||||
"""Search a Peer"""
|
||||
stmt = await crud.search(query, workspace_name=workspace_id, peer_name=peer_id)
|
||||
|
||||
return await paginate(db, stmt)
|
||||
|
|
@ -1,14 +1,11 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from anthropic import AsyncMessageStreamManager
|
||||
from fastapi import APIRouter, Body, Depends, Path, Query
|
||||
from fastapi.exceptions import HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi import APIRouter, Body, Depends, Path, Query, Response
|
||||
from fastapi_pagination import Page
|
||||
from fastapi_pagination.ext.sqlalchemy import paginate
|
||||
|
||||
from src import agent, crud, schemas
|
||||
from src import crud, schemas
|
||||
from src.dependencies import db
|
||||
from src.exceptions import (
|
||||
AuthenticationException,
|
||||
|
|
@ -16,140 +13,106 @@ from src.exceptions import (
|
|||
ValidationException,
|
||||
)
|
||||
from src.security import JWTParams, require_auth
|
||||
from src.utils import history
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps/{app_id}/users/{user_id}/sessions",
|
||||
prefix="/workspaces/{workspace_id}/sessions",
|
||||
tags=["sessions"],
|
||||
)
|
||||
|
||||
|
||||
@router.get(
|
||||
@router.post(
|
||||
"",
|
||||
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"
|
||||
async def get_or_create_session(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session: schemas.SessionCreate = Body(
|
||||
..., description="Session creation parameters"
|
||||
),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
db=db,
|
||||
):
|
||||
"""
|
||||
Get a specific session for a user.
|
||||
Get a specific session in a workspace.
|
||||
|
||||
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.
|
||||
If session_id is provided as a query parameter, it verifies the session is in the workspace.
|
||||
Otherwise, it uses the session_id from the JWT token for verification.
|
||||
"""
|
||||
# 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.w is not None and jwt_params.w != workspace_id:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
|
||||
# Use session from JWT if not provided in query
|
||||
if session.name:
|
||||
if (
|
||||
not jwt_params.ad
|
||||
and jwt_params.se is not None
|
||||
and jwt_params.se != session_id
|
||||
and jwt_params.s is not None
|
||||
and jwt_params.s != session.name
|
||||
):
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
target_session_id = session_id
|
||||
else:
|
||||
# Use session_id from JWT
|
||||
if not jwt_params.se:
|
||||
if not jwt_params.s:
|
||||
raise AuthenticationException(
|
||||
"Session ID not found in query parameter or JWT"
|
||||
)
|
||||
target_session_id = jwt_params.se
|
||||
session.name = jwt_params.s
|
||||
|
||||
# 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
|
||||
)
|
||||
# Handle session creation with proper error handling
|
||||
try:
|
||||
return await crud.get_or_create_session(
|
||||
db, workspace_name=workspace_id, session=session
|
||||
)
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to get or create session {session.name}: {str(e)}")
|
||||
raise ValidationException(str(e)) from e
|
||||
|
||||
|
||||
@router.post(
|
||||
"/list",
|
||||
response_model=Page[schemas.Session],
|
||||
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
|
||||
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
|
||||
)
|
||||
async def get_sessions(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
options: Optional[schemas.SessionGet] = Body(
|
||||
None, 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"""
|
||||
"""Get All Sessions in a Workspace"""
|
||||
filter_param = None
|
||||
is_active_param = False # Default to None, meaning no filter on is_active
|
||||
is_active_param = False # Default from schema
|
||||
|
||||
if options:
|
||||
if hasattr(options, 'filter') and options.filter:
|
||||
if hasattr(options, "filter") and options.filter:
|
||||
filter_param = options.filter
|
||||
if filter_param == {}: # Explicitly check for empty dict
|
||||
if filter_param == {}: # Explicitly check for empty dict
|
||||
filter_param = None
|
||||
if hasattr(options, 'is_active'): # Check if is_active is present
|
||||
if hasattr(options, "is_active"): # Check if is_active is present
|
||||
is_active_param = options.is_active
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_sessions(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
reverse=reverse,
|
||||
workspace_name=workspace_id,
|
||||
is_active=is_active_param,
|
||||
filter=filter_param,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@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 = 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"""
|
||||
try:
|
||||
session_obj = await crud.create_session(
|
||||
db, app_id=app_id, user_id=user_id, session=session
|
||||
)
|
||||
logger.info(f"Session created successfully for user {user_id}")
|
||||
return session_obj
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to create session: {str(e)}")
|
||||
raise ValidationException(str(e)) from e
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{session_id}",
|
||||
response_model=schemas.Session,
|
||||
dependencies=[
|
||||
Depends(
|
||||
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
|
||||
)
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def update_session(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session to update"),
|
||||
session: schemas.SessionUpdate = Body(
|
||||
..., description="Updated session parameters"
|
||||
|
|
@ -159,7 +122,7 @@ async def update_session(
|
|||
"""Update the metadata of a Session"""
|
||||
try:
|
||||
updated_session = await crud.update_session(
|
||||
db, app_id=app_id, user_id=user_id, session_id=session_id, session=session
|
||||
db, workspace_name=workspace_id, session_name=session_id, session=session
|
||||
)
|
||||
logger.info(f"Session {session_id} updated successfully")
|
||||
return updated_session
|
||||
|
|
@ -171,21 +134,18 @@ async def update_session(
|
|||
@router.delete(
|
||||
"/{session_id}",
|
||||
dependencies=[
|
||||
Depends(
|
||||
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
|
||||
)
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def delete_session(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session to delete"),
|
||||
db=db,
|
||||
):
|
||||
"""Delete a session by marking it as inactive"""
|
||||
try:
|
||||
await crud.delete_session(
|
||||
db, app_id=app_id, user_id=user_id, session_id=session_id
|
||||
db, workspace_name=workspace_id, session_name=session_id
|
||||
)
|
||||
logger.info(f"Session {session_id} deleted successfully")
|
||||
return {"message": "Session deleted successfully"}
|
||||
|
|
@ -194,93 +154,322 @@ async def delete_session(
|
|||
raise ResourceNotFoundException("Session not found") from e
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{session_id}/chat",
|
||||
response_model=schemas.DialecticResponse,
|
||||
responses={
|
||||
200: {
|
||||
"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 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"
|
||||
),
|
||||
):
|
||||
|
||||
"""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,
|
||||
queries=options.queries,
|
||||
)
|
||||
else:
|
||||
|
||||
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,
|
||||
dependencies=[
|
||||
Depends(
|
||||
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
|
||||
)
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def clone_session(
|
||||
app_id: str = Path(..., description="ID of the app"),
|
||||
user_id: str = Path(..., description="ID of the user"),
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session to clone"),
|
||||
db=db,
|
||||
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, optionally up to a specific message"""
|
||||
try:
|
||||
# TODO: Update crud.clone_session to work with new paradigm
|
||||
cloned_session = await crud.clone_session(
|
||||
db,
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
original_session_id=session_id,
|
||||
workspace_name=workspace_id,
|
||||
original_session_name=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
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{session_id}/peers",
|
||||
response_model=schemas.Session,
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def add_peers_to_session(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
peers: dict[str, schemas.SessionPeerConfig] = Body(
|
||||
..., description="List of peer IDs to add to the session"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Add peers to a session"""
|
||||
try:
|
||||
workspace_name, session_name = workspace_id, session_id
|
||||
session = await crud.get_or_create_session(
|
||||
db,
|
||||
session=schemas.SessionCreate(
|
||||
name=session_name,
|
||||
peers=peers,
|
||||
),
|
||||
workspace_name=workspace_name,
|
||||
)
|
||||
logger.info(f"Added peers to session {session_name} successfully")
|
||||
return session
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to add peers to session {session_name}: {str(e)}")
|
||||
raise ResourceNotFoundException("Session not found") from e
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{session_id}/peers",
|
||||
response_model=schemas.Session,
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def set_session_peers(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
peers: dict[str, schemas.SessionPeerConfig] = Body(
|
||||
..., description="List of peer IDs to set for the session"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Set the peers in a session"""
|
||||
try:
|
||||
workspace_name, session_name = workspace_id, session_id
|
||||
await crud.set_peers_for_session(
|
||||
db,
|
||||
workspace_name=workspace_name,
|
||||
session_name=session_name,
|
||||
peer_names=peers,
|
||||
)
|
||||
# Get the session to return
|
||||
session = await crud.get_or_create_session(
|
||||
db,
|
||||
session=schemas.SessionCreate(name=session_name),
|
||||
workspace_name=workspace_name,
|
||||
)
|
||||
logger.info(f"Set peers for session {session_name} successfully")
|
||||
return session
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to set peers for session {session_name}: {str(e)}")
|
||||
raise ResourceNotFoundException("Failed to set peers for session") from e
|
||||
|
||||
|
||||
@router.delete(
|
||||
"/{session_id}/peers",
|
||||
response_model=schemas.Session,
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def remove_peers_from_session(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
peers: list[str] = Body(
|
||||
..., description="List of peer IDs to remove from the session"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Remove peers from a session"""
|
||||
try:
|
||||
session_name = session_id
|
||||
await crud.remove_peers_from_session(
|
||||
db,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_name,
|
||||
peer_names=set(peers),
|
||||
)
|
||||
# Get the session to return
|
||||
session = await crud.get_or_create_session(
|
||||
db,
|
||||
session=schemas.SessionCreate(name=session_name),
|
||||
workspace_name=workspace_id,
|
||||
)
|
||||
logger.info(f"Removed peers from session {session_name} successfully")
|
||||
return session
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to remove peers from session {session_name}: {str(e)}")
|
||||
raise ResourceNotFoundException("Session not found") from e
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{session_id}/peers/{peer_id}/config",
|
||||
response_model=schemas.SessionPeerConfig,
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def get_peer_config(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
db=db,
|
||||
):
|
||||
"""Get the configuration for a peer in a session"""
|
||||
return await crud.get_peer_config(
|
||||
db,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
peer_id=peer_id,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{session_id}/peers/{peer_id}/config",
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def set_peer_config(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
peer_id: str = Path(..., description="ID of the peer"),
|
||||
config: schemas.SessionPeerConfig = Body(..., description="Peer configuration"),
|
||||
db=db,
|
||||
):
|
||||
"""Set the configuration for a peer in a session"""
|
||||
try:
|
||||
await crud.set_peer_config(
|
||||
db,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
peer_id=peer_id,
|
||||
config=config,
|
||||
)
|
||||
logger.info(
|
||||
f"Set peer config for {peer_id} in session {session_id} successfully"
|
||||
)
|
||||
return Response(status_code=200)
|
||||
except ValueError as e:
|
||||
logger.warning(
|
||||
f"Failed to set peer config for {peer_id} in session {session_id}: {str(e)}"
|
||||
)
|
||||
raise ResourceNotFoundException("Session not found") from e
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{session_id}/peers",
|
||||
response_model=Page[schemas.Peer],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def get_session_peers(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
db=db,
|
||||
):
|
||||
"""Get peers from a session"""
|
||||
try:
|
||||
session_name = session_id
|
||||
peers_query = await crud.get_peers_from_session(
|
||||
workspace_name=workspace_id, session_name=session_name
|
||||
)
|
||||
return await paginate(db, peers_query)
|
||||
except ValueError as e:
|
||||
logger.warning(f"Failed to get peers from session {session_name}: {str(e)}")
|
||||
raise ResourceNotFoundException("Session not found") from e
|
||||
|
||||
|
||||
@router.get(
|
||||
"/{session_id}/context",
|
||||
response_model=schemas.SessionContext,
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def get_session_context(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
tokens: Optional[int] = Query(
|
||||
None,
|
||||
description="Number of tokens to use for the context. Includes summary if set to true",
|
||||
),
|
||||
summary: bool = Query(
|
||||
False,
|
||||
description="Whether to summarize the session history prior to the cutoff message",
|
||||
), # default to false
|
||||
db=db,
|
||||
):
|
||||
"""
|
||||
Produce a context object from the session. The caller provides a token limit which the entire context must fit into.
|
||||
To do this, we allocate 40% of the token limit to the summary, and 60% to recent messages -- as many as can fit.
|
||||
If the caller does not want a summary, we allocate all the tokens to recent messages.
|
||||
The default token limit if not provided is 2048. (TODO: make this configurable)
|
||||
"""
|
||||
token_limit = tokens or 2048
|
||||
summary_tokens = int(token_limit * 0.4) if summary else 0
|
||||
messages_tokens = token_limit - summary_tokens
|
||||
|
||||
# Get the messages to return verbatim
|
||||
messages_stmt = await crud.get_messages(
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
token_limit=messages_tokens,
|
||||
)
|
||||
result = await db.execute(messages_stmt)
|
||||
messages = list(result.scalars().all())
|
||||
|
||||
# Get the most recently created summary for the session
|
||||
last_summary = await history.get_summary(
|
||||
db,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
)
|
||||
|
||||
# Get messages between the last summary and the first message we'll return verbatim, if any
|
||||
messages_before = await crud.get_messages_id_range(
|
||||
db,
|
||||
workspace_name=workspace_id,
|
||||
session_name=session_id,
|
||||
peer_name=None,
|
||||
start_id=last_summary["message_id"] if last_summary else 0,
|
||||
end_id=messages[0].id if messages else None,
|
||||
)
|
||||
|
||||
# Make a summary if the user wants one
|
||||
if summary_tokens > 0:
|
||||
# Make a *new* summary if there are unsummarized messages between the last summary and the ones
|
||||
# we'll return verbatim, or if the last summary is too many tokens -- otherwise, just use the last summary
|
||||
if (
|
||||
not last_summary
|
||||
or len(messages_before) > 0
|
||||
or last_summary["token_count"] > summary_tokens
|
||||
):
|
||||
new_summary = await history.create_summary(
|
||||
messages=messages_before,
|
||||
max_tokens=summary_tokens,
|
||||
)
|
||||
summary_content = new_summary["content"]
|
||||
else:
|
||||
summary_content = last_summary["content"]
|
||||
summary_tokens = last_summary["token_count"]
|
||||
else:
|
||||
summary_content = ""
|
||||
|
||||
return schemas.SessionContext(
|
||||
name=session_id,
|
||||
messages=messages,
|
||||
summary=summary_content,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{session_id}/search",
|
||||
response_model=Page[schemas.Message],
|
||||
dependencies=[
|
||||
Depends(require_auth(workspace_name="workspace_id", session_name="session_id"))
|
||||
],
|
||||
)
|
||||
async def search_session(
|
||||
workspace_id: str = Path(..., description="ID of the workspace"),
|
||||
session_id: str = Path(..., description="ID of the session"),
|
||||
query: str = Body(..., description="Search query"),
|
||||
db=db,
|
||||
):
|
||||
"""Search a Session"""
|
||||
stmt = await crud.search(
|
||||
query, workspace_name=workspace_id, session_name=session_id
|
||||
)
|
||||
|
||||
return await paginate(db, stmt)
|
||||
|
|
|
|||
|
|
@ -1,162 +0,0 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
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 JWTParams, require_auth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/apps/{app_id}/users",
|
||||
tags=["users"],
|
||||
)
|
||||
|
||||
|
||||
@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 = Path(..., description="ID of the app"),
|
||||
user: schemas.UserCreate = Body(..., description="User creation parameters"),
|
||||
db=db,
|
||||
):
|
||||
"""Create a new User"""
|
||||
user_obj = await crud.create_user(db, app_id=app_id, user=user)
|
||||
return user_obj
|
||||
|
||||
|
||||
@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 = Path(..., description="ID of the app"),
|
||||
options: Optional[schemas.UserGet] = Body(
|
||||
None, 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"""
|
||||
filter_param = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}:
|
||||
filter_param = None
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_users(app_id=app_id, reverse=reverse, filter=filter_param),
|
||||
)
|
||||
|
||||
|
||||
@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 = Path(..., description="ID of the app"),
|
||||
name: str = Path(..., description="Name of the user to retrieve"),
|
||||
db=db,
|
||||
):
|
||||
"""Get a User by name"""
|
||||
user = await crud.get_user_by_name(db, app_id=app_id, name=name)
|
||||
return user
|
||||
|
||||
|
||||
@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 or create a new one by the input name"""
|
||||
try:
|
||||
user = await crud.get_user_by_name(db, app_id=app_id, name=name)
|
||||
return user
|
||||
except ResourceNotFoundException:
|
||||
# User doesn't exist, create it
|
||||
user = await create_user(
|
||||
db=db, app_id=app_id, user=schemas.UserCreate(name=name)
|
||||
)
|
||||
return 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 = 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"""
|
||||
updated_user = await crud.update_user(db, app_id=app_id, user_id=user_id, user=user)
|
||||
return updated_user
|
||||
|
|
@ -0,0 +1,107 @@
|
|||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Body, Depends, Path
|
||||
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
|
||||
from src.security import JWTParams, require_auth
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(
|
||||
prefix="/workspaces",
|
||||
tags=["workspaces"],
|
||||
)
|
||||
|
||||
|
||||
@router.post("", response_model=schemas.Workspace)
|
||||
async def get_or_create_workspace(
|
||||
workspace: schemas.WorkspaceCreate = Body(
|
||||
..., description="Workspace creation parameters"
|
||||
),
|
||||
jwt_params: JWTParams = Depends(require_auth()),
|
||||
db=db,
|
||||
):
|
||||
"""
|
||||
Get a Workspace by ID.
|
||||
|
||||
If workspace_id is provided as a query parameter, it uses that (must match JWT workspace_id).
|
||||
Otherwise, it uses the workspace_id from the JWT token.
|
||||
"""
|
||||
# If workspace_id provided in query, check if it matches jwt or user is admin
|
||||
if workspace.name:
|
||||
if not jwt_params.ad and jwt_params.w != workspace.name:
|
||||
raise AuthenticationException("Unauthorized access to resource")
|
||||
else:
|
||||
# Use workspace_id from JWT
|
||||
if not jwt_params.w:
|
||||
raise AuthenticationException(
|
||||
"Workspace ID not found in query parameter or JWT"
|
||||
)
|
||||
workspace.name = jwt_params.w
|
||||
|
||||
return await crud.get_or_create_workspace(db, workspace=workspace)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/list",
|
||||
response_model=Page[schemas.Workspace],
|
||||
dependencies=[Depends(require_auth(admin=True))],
|
||||
)
|
||||
async def get_all_workspaces(
|
||||
options: Optional[schemas.WorkspaceGet] = Body(
|
||||
None, description="Filtering and pagination options for the workspaces list"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Get all Workspaces"""
|
||||
filter_param = None
|
||||
if options and hasattr(options, "filter"):
|
||||
filter_param = options.filter
|
||||
if filter_param == {}:
|
||||
filter_param = None
|
||||
|
||||
return await paginate(
|
||||
db,
|
||||
await crud.get_all_workspaces(filter=filter_param),
|
||||
)
|
||||
|
||||
|
||||
@router.put(
|
||||
"/{workspace_id}",
|
||||
response_model=schemas.Workspace,
|
||||
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
|
||||
)
|
||||
async def update_workspace(
|
||||
workspace_id: str = Path(..., description="ID of the workspace to update"),
|
||||
workspace: schemas.WorkspaceUpdate = Body(
|
||||
..., description="Updated workspace parameters"
|
||||
),
|
||||
db=db,
|
||||
):
|
||||
"""Update a Workspace"""
|
||||
# ResourceNotFoundException will be caught by global handler if workspace not found
|
||||
honcho_workspace = await crud.update_workspace(
|
||||
db, workspace_name=workspace_id, workspace=workspace
|
||||
)
|
||||
return honcho_workspace
|
||||
|
||||
|
||||
@router.post(
|
||||
"/{workspace_id}/search",
|
||||
response_model=Page[schemas.Message],
|
||||
dependencies=[Depends(require_auth(workspace_name="workspace_id"))],
|
||||
)
|
||||
async def search_workspace(
|
||||
workspace_id: str = Path(..., description="ID of the workspace to search"),
|
||||
query: str = Body(..., description="Search query"),
|
||||
db=db,
|
||||
):
|
||||
"""Search a Workspace"""
|
||||
stmt = await crud.search(query, workspace_name=workspace_id)
|
||||
|
||||
return await paginate(db, stmt)
|
||||
321
src/schemas.py
321
src/schemas.py
|
|
@ -1,67 +1,92 @@
|
|||
import datetime
|
||||
from typing import Annotated
|
||||
from typing import Annotated, Optional
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, computed_field, field_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
RESOURCE_NAME_PATTERN = r"^[a-zA-Z0-9_-]+$"
|
||||
|
||||
|
||||
class AppBase(BaseModel):
|
||||
class WorkspaceBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class AppCreate(AppBase):
|
||||
name: Annotated[str, Field(min_length=1, max_length=100)]
|
||||
class WorkspaceCreate(WorkspaceBase):
|
||||
name: Annotated[
|
||||
str,
|
||||
Field(alias="id", min_length=1, max_length=100, pattern=RESOURCE_NAME_PATTERN),
|
||||
]
|
||||
metadata: dict = {}
|
||||
configuration: dict = {}
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class AppGet(AppBase):
|
||||
class WorkspaceGet(WorkspaceBase):
|
||||
filter: dict | None = None
|
||||
|
||||
|
||||
class AppUpdate(AppBase):
|
||||
name: str | None = None
|
||||
class WorkspaceUpdate(WorkspaceBase):
|
||||
metadata: dict | None = None
|
||||
configuration: dict | None = None
|
||||
|
||||
|
||||
class App(AppBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
name: str
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
class Workspace(WorkspaceBase):
|
||||
name: str = Field(serialization_alias="id")
|
||||
h_metadata: dict = Field(default_factory=dict, serialization_alias="metadata")
|
||||
configuration: dict = Field(default_factory=dict)
|
||||
created_at: datetime.datetime
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
||||
|
||||
|
||||
class PeerBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class PeerCreate(PeerBase):
|
||||
name: Annotated[
|
||||
str,
|
||||
Field(alias="id", min_length=1, max_length=100, pattern=RESOURCE_NAME_PATTERN),
|
||||
]
|
||||
metadata: dict | None = None
|
||||
configuration: dict | None = None
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class PeerGet(PeerBase):
|
||||
filter: dict | None = None
|
||||
|
||||
|
||||
class PeerUpdate(PeerBase):
|
||||
metadata: dict | None = None
|
||||
configuration: dict | None = None
|
||||
|
||||
|
||||
class Peer(PeerBase):
|
||||
name: str = Field(serialization_alias="id")
|
||||
workspace_name: str = Field(serialization_alias="workspace_id")
|
||||
created_at: datetime.datetime
|
||||
h_metadata: dict = Field(default_factory=dict, serialization_alias="metadata")
|
||||
configuration: dict = Field(default_factory=dict)
|
||||
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
||||
|
||||
|
||||
class PeerRepresentationGet(BaseModel):
|
||||
session_id: str = Field(
|
||||
..., description="Get the working representation within this session"
|
||||
)
|
||||
target: Optional[str] = Field(
|
||||
None,
|
||||
description="Optional peer ID to get the representation for, from the perspective of this peer",
|
||||
)
|
||||
|
||||
|
||||
class UserBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class UserCreate(UserBase):
|
||||
name: Annotated[str, Field(min_length=1, max_length=100)]
|
||||
metadata: dict = {}
|
||||
|
||||
|
||||
class UserGet(UserBase):
|
||||
filter: dict | None = None
|
||||
|
||||
|
||||
class UserUpdate(UserBase):
|
||||
name: str | None = None
|
||||
metadata: dict | None = None
|
||||
|
||||
|
||||
class User(UserBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
name: str
|
||||
app_id: str
|
||||
created_at: datetime.datetime
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
class PeerConfig(BaseModel):
|
||||
observe_me: bool = Field(
|
||||
default=True,
|
||||
description="Whether honcho should form a global theory-of-mind representation of this peer",
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -71,8 +96,8 @@ class MessageBase(BaseModel):
|
|||
|
||||
class MessageCreate(MessageBase):
|
||||
content: Annotated[str, Field(min_length=0, max_length=50000)]
|
||||
is_user: bool
|
||||
metadata: dict = {}
|
||||
peer_name: str = Field(alias="peer_id")
|
||||
metadata: dict | None = None
|
||||
|
||||
|
||||
class MessageGet(MessageBase):
|
||||
|
|
@ -80,31 +105,53 @@ class MessageGet(MessageBase):
|
|||
|
||||
|
||||
class MessageUpdate(MessageBase):
|
||||
metadata: dict
|
||||
metadata: dict | None = None
|
||||
|
||||
|
||||
class Message(MessageBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
public_id: str = Field(serialization_alias="id")
|
||||
content: str
|
||||
is_user: bool
|
||||
session_id: str
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
peer_name: str = Field(serialization_alias="peer_id")
|
||||
session_name: str | None = Field(serialization_alias="session_id")
|
||||
h_metadata: dict = Field(default_factory=dict, serialization_alias="metadata")
|
||||
created_at: datetime.datetime
|
||||
app_id: str
|
||||
user_id: str
|
||||
workspace_name: str = Field(serialization_alias="workspace_id")
|
||||
token_count: int
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
)
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
||||
|
||||
|
||||
class MessageBatchCreate(BaseModel):
|
||||
"""Schema for batch message creation with a max of 100 messages"""
|
||||
|
||||
messages: list[MessageCreate] = Field(..., min_length=1, max_length=100)
|
||||
|
||||
|
||||
class SessionBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class SessionPeerConfig(BaseModel):
|
||||
observe_others: bool = Field(
|
||||
default=False,
|
||||
description="Whether this peer should form a session-level theory-of-mind representation of other peers in the session",
|
||||
)
|
||||
observe_me: bool = Field(
|
||||
default=True,
|
||||
description="Whether other peers in this session should try to form a session-level theory-of-mind representation of this peer",
|
||||
)
|
||||
|
||||
|
||||
class SessionCreate(SessionBase):
|
||||
metadata: dict = {}
|
||||
name: Annotated[
|
||||
str,
|
||||
Field(alias="id", min_length=1, max_length=100, pattern=RESOURCE_NAME_PATTERN),
|
||||
]
|
||||
metadata: dict | None = None
|
||||
peer_names: dict[str, SessionPeerConfig] | None = Field(default=None, alias="peers")
|
||||
configuration: dict | None = None
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class SessionGet(SessionBase):
|
||||
|
|
@ -113,120 +160,27 @@ class SessionGet(SessionBase):
|
|||
|
||||
|
||||
class SessionUpdate(SessionBase):
|
||||
metadata: dict
|
||||
metadata: dict | None = None
|
||||
configuration: dict | None = None
|
||||
|
||||
|
||||
class Session(SessionBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
name: str = Field(serialization_alias="id")
|
||||
is_active: bool
|
||||
user_id: str
|
||||
app_id: str
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
workspace_name: str = Field(serialization_alias="workspace_id")
|
||||
h_metadata: dict = Field(default_factory=dict, serialization_alias="metadata")
|
||||
configuration: dict = Field(default_factory=dict)
|
||||
created_at: datetime.datetime
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
)
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
||||
|
||||
|
||||
class MetamessageBase(BaseModel):
|
||||
pass
|
||||
class SessionContext(SessionBase):
|
||||
name: str = Field(serialization_alias="id")
|
||||
messages: list[Message]
|
||||
summary: str
|
||||
|
||||
|
||||
class MetamessageCreate(MetamessageBase):
|
||||
label: Annotated[str, Field(min_length=1, max_length=50, alias='metamessage_type')]
|
||||
content: Annotated[str, Field(min_length=0, max_length=50000)]
|
||||
session_id: str | None = None
|
||||
message_id: str | None = None
|
||||
metadata: dict = {}
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class MetamessageGet(MetamessageBase):
|
||||
label: str | None = Field(default=None, alias='metamessage_type')
|
||||
session_id: str | None = None
|
||||
message_id: str | None = None
|
||||
filter: dict | None = None
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class MetamessageUpdate(MetamessageBase):
|
||||
session_id: str | None = None
|
||||
message_id: str | None = None
|
||||
label: str | None = Field(default=None, alias='metamessage_type')
|
||||
metadata: dict | None = None
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
|
||||
class Metamessage(MetamessageBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
label: str
|
||||
content: str
|
||||
user_id: str
|
||||
app_id: str
|
||||
session_id: str | None
|
||||
message_id: str | None
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
created_at: datetime.datetime
|
||||
|
||||
# Included for backwards compatibility with the old metamessage_type field
|
||||
@computed_field
|
||||
@property
|
||||
def metamessage_type(self) -> str:
|
||||
return self.label
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
)
|
||||
|
||||
|
||||
class CollectionBase(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
class CollectionCreate(CollectionBase):
|
||||
name: Annotated[str, Field(min_length=1, max_length=100)]
|
||||
metadata: dict = {}
|
||||
|
||||
@field_validator("name")
|
||||
def validate_name(cls, v):
|
||||
if v.lower() == "honcho":
|
||||
raise ValueError("Collection name cannot be 'honcho'")
|
||||
return v
|
||||
|
||||
|
||||
class CollectionGet(CollectionBase):
|
||||
filter: dict | None = None
|
||||
|
||||
|
||||
class CollectionUpdate(CollectionBase):
|
||||
name: str | None = None
|
||||
metadata: dict | None = None
|
||||
|
||||
@field_validator("name")
|
||||
def validate_name(cls, v):
|
||||
if v is not None and v.lower() == "honcho":
|
||||
raise ValueError("Collection name cannot be 'honcho'")
|
||||
return v
|
||||
|
||||
|
||||
class Collection(CollectionBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
name: str
|
||||
user_id: str
|
||||
app_id: str
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
created_at: datetime.datetime
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
)
|
||||
model_config = ConfigDict(from_attributes=True, populate_by_name=True)
|
||||
|
||||
|
||||
class DocumentBase(BaseModel):
|
||||
|
|
@ -238,37 +192,14 @@ class DocumentCreate(DocumentBase):
|
|||
metadata: dict = {}
|
||||
|
||||
|
||||
class DocumentGet(DocumentBase):
|
||||
filter: dict | None = None
|
||||
|
||||
|
||||
class DocumentQuery(DocumentBase):
|
||||
query: Annotated[str, Field(min_length=1, max_length=1000)]
|
||||
filter: dict | None = None
|
||||
top_k: int = Field(default=5, ge=1, le=50)
|
||||
|
||||
|
||||
class DocumentUpdate(DocumentBase):
|
||||
metadata: dict | None = Field(None, max_length=10000)
|
||||
content: Annotated[str | None, Field(min_length=1, max_length=100000)] = None
|
||||
|
||||
|
||||
class Document(DocumentBase):
|
||||
public_id: str = Field(serialization_alias='id')
|
||||
content: str
|
||||
h_metadata: dict = Field(default={}, serialization_alias='metadata')
|
||||
created_at: datetime.datetime
|
||||
collection_id: str
|
||||
app_id: str
|
||||
user_id: str
|
||||
|
||||
model_config = ConfigDict(
|
||||
from_attributes=True,
|
||||
populate_by_name=True
|
||||
)
|
||||
|
||||
|
||||
class DialecticOptions(BaseModel):
|
||||
session_id: Optional[str] = Field(
|
||||
None, description="ID of the session to scope the representation to"
|
||||
)
|
||||
target: Optional[str] = Field(
|
||||
None,
|
||||
description="Optional peer to get the representation for, from the perspective of this peer",
|
||||
)
|
||||
queries: str | list[str]
|
||||
stream: bool = False
|
||||
|
||||
|
|
@ -289,9 +220,3 @@ class DialecticOptions(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)
|
||||
|
|
|
|||
108
src/security.py
108
src/security.py
|
|
@ -41,10 +41,14 @@ security = HTTPBearer(
|
|||
class JWTParams(BaseModel):
|
||||
"""
|
||||
JWT parameters used to produce tokens valid for different routes.
|
||||
Hierarchy: app > user > (session / collection)
|
||||
Workspaces are the top level of the hierarchy -- a workspace key will
|
||||
give access to all peers/sessions/collections in that workspace.
|
||||
|
||||
All routers require at least the most tightly scoped parameter.
|
||||
Routes will accept a JWT with a scope higher in the hierarchy.
|
||||
A session key will allow the listing and creation of messages in
|
||||
that session.
|
||||
|
||||
A peer key will allow the listing and creation of peer-level messages
|
||||
and querying the peer's dialectic endpoint.
|
||||
|
||||
Names shortened to minimize token size. Timestamp is included
|
||||
so that many unique tokens can be generated for the same resource.
|
||||
|
|
@ -56,19 +60,17 @@ class JWTParams(BaseModel):
|
|||
`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
|
||||
`w`: (string) workspace name
|
||||
`p`: (string) peer name
|
||||
`s`: (string) session name
|
||||
"""
|
||||
|
||||
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
|
||||
w: Optional[str] = None
|
||||
p: Optional[str] = None
|
||||
s: Optional[str] = None
|
||||
|
||||
|
||||
def create_admin_jwt() -> str:
|
||||
|
|
@ -104,14 +106,12 @@ async def verify_jwt(token: str) -> JWTParams:
|
|||
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"]
|
||||
if "w" in decoded:
|
||||
params.w = decoded["w"]
|
||||
if "p" in decoded:
|
||||
params.p = decoded["p"]
|
||||
if "s" in decoded:
|
||||
params.s = decoded["s"]
|
||||
return params
|
||||
except jwt.PyJWTError:
|
||||
raise AuthenticationException("Invalid JWT") from None
|
||||
|
|
@ -119,10 +119,9 @@ async def verify_jwt(token: str) -> JWTParams:
|
|||
|
||||
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,
|
||||
workspace_name: Optional[str] = None,
|
||||
peer_name: Optional[str] = None,
|
||||
session_name: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Generate a dependency that requires authentication for the given parameters.
|
||||
|
|
@ -133,35 +132,30 @@ def require_auth(
|
|||
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
|
||||
workspace_name_param = (
|
||||
request.path_params.get(workspace_name)
|
||||
or request.query_params.get(workspace_name)
|
||||
if workspace_name
|
||||
else None
|
||||
)
|
||||
user_id_param = (
|
||||
request.path_params.get(user_id) or request.query_params.get(user_id)
|
||||
if user_id
|
||||
peer_name_param = (
|
||||
request.path_params.get(peer_name) or request.query_params.get(peer_name)
|
||||
if peer_name
|
||||
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
|
||||
session_name_param = (
|
||||
request.path_params.get(session_name)
|
||||
or request.query_params.get(session_name)
|
||||
if session_name
|
||||
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,
|
||||
workspace_name=workspace_name_param,
|
||||
peer_name=peer_name_param,
|
||||
session_name=session_name_param,
|
||||
)
|
||||
|
||||
return auth_dependency
|
||||
|
|
@ -170,10 +164,9 @@ def require_auth(
|
|||
async def auth(
|
||||
credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)],
|
||||
admin: Optional[bool] = None,
|
||||
app_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
session_id: Optional[str] = None,
|
||||
collection_id: Optional[str] = None,
|
||||
workspace_name: Optional[str] = None,
|
||||
peer_name: Optional[str] = None,
|
||||
session_name: Optional[str] = None,
|
||||
) -> JWTParams:
|
||||
"""Authenticate the given JWT and return the decoded parameters."""
|
||||
if not USE_AUTH:
|
||||
|
|
@ -190,24 +183,23 @@ async def auth(
|
|||
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:
|
||||
# For session level access
|
||||
if session_name and jwt_params.s == session_name:
|
||||
if workspace_name and jwt_params.w != workspace_name:
|
||||
raise AuthenticationException("JWT not permissioned for this resource")
|
||||
return jwt_params
|
||||
|
||||
# For user level access - can access all sessions/collections under this user
|
||||
if user_id and jwt_params.us == user_id:
|
||||
# For peer level access
|
||||
if peer_name and jwt_params.p == peer_name:
|
||||
if workspace_name and jwt_params.w != workspace_name:
|
||||
raise AuthenticationException("JWT not permissioned for this resource")
|
||||
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:
|
||||
# For workspace level access - can access all peers/sessions under this workspace
|
||||
if workspace_name and jwt_params.w == workspace_name:
|
||||
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)
|
||||
if any([session_name, peer_name, workspace_name]):
|
||||
raise AuthenticationException("JWT not permissioned for this resource")
|
||||
|
||||
# Route did not specify any parameters, so it should parse parameters itself
|
||||
|
|
|
|||
|
|
@ -1,28 +1,48 @@
|
|||
import datetime
|
||||
import logging
|
||||
from enum import Enum
|
||||
from typing import Optional, Union, cast
|
||||
from typing import Optional, TypedDict
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src.utils.model_client import ModelClient, ModelProvider
|
||||
|
||||
from .. import models
|
||||
from .. import crud, models
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# TypedDict definitions for summary data
|
||||
class Summary(TypedDict):
|
||||
"""
|
||||
A summary object. Stored in session metadata and used in a session's get_context.
|
||||
|
||||
Attributes:
|
||||
content: The summary text.
|
||||
message_count: The number of messages covered by this summary.
|
||||
summary_type: The type of summary (short or long).
|
||||
created_at: The timestamp of when the summary was created (ISO format string).
|
||||
message_id: The primary key ID of the message that triggered this summary.
|
||||
token_count: The number of tokens in the summary text.
|
||||
"""
|
||||
|
||||
content: str
|
||||
message_count: int
|
||||
summary_type: str
|
||||
created_at: str
|
||||
message_id: int
|
||||
token_count: int
|
||||
|
||||
|
||||
# Export the public functions
|
||||
__all__ = [
|
||||
"get_session_summaries",
|
||||
"get_messages_since_message",
|
||||
"get_messages_since_latest_summary",
|
||||
"get_summary",
|
||||
"create_summary",
|
||||
"save_summary_metamessage",
|
||||
"save_summary",
|
||||
"get_summarized_history",
|
||||
"should_create_summary",
|
||||
"MESSAGES_PER_SHORT_SUMMARY",
|
||||
"MESSAGES_PER_LONG_SUMMARY",
|
||||
"SummaryType",
|
||||
"Summary",
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -31,7 +51,7 @@ 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
|
||||
# The types of summary to store in the session metadata
|
||||
class SummaryType(Enum):
|
||||
SHORT = "honcho_chat_summary_short"
|
||||
LONG = "honcho_chat_summary_long"
|
||||
|
|
@ -42,110 +62,64 @@ DEFAULT_PROVIDER = ModelProvider.GEMINI
|
|||
DEFAULT_MODEL = "gemini-2.0-flash-lite"
|
||||
|
||||
|
||||
async def get_session_summaries(
|
||||
async def get_summary(
|
||||
db: AsyncSession,
|
||||
session_id: str,
|
||||
workspace_name: str,
|
||||
session_name: str,
|
||||
summary_type: SummaryType = SummaryType.SHORT,
|
||||
only_latest: bool = False,
|
||||
) -> Union[list[models.Metamessage], Optional[models.Metamessage]]:
|
||||
) -> Optional[Summary]:
|
||||
"""
|
||||
Get summaries for a given session.
|
||||
Get summary for a given session or peer.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
session_id: The session ID
|
||||
workspace_name: The workspace name
|
||||
session_name: The session name
|
||||
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
|
||||
The summary data dictionary, or None if no summary exists
|
||||
"""
|
||||
# Determine the metamessage type based on summary_type
|
||||
from src.exceptions import ResourceNotFoundException
|
||||
|
||||
label = (
|
||||
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.label == label)
|
||||
.order_by(models.Metamessage.id.desc())
|
||||
)
|
||||
try:
|
||||
session = await crud.get_session(db, session_name, workspace_name)
|
||||
except ResourceNotFoundException:
|
||||
# If session doesn't exist, there's no summary to retrieve
|
||||
return None
|
||||
|
||||
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())
|
||||
summaries = session.internal_metadata.get("summaries", {})
|
||||
if not summaries or label not in summaries:
|
||||
return None
|
||||
return summaries[label]
|
||||
|
||||
|
||||
async def create_summary(
|
||||
messages: list[models.Message],
|
||||
previous_summary: Optional[str] = None,
|
||||
previous_summary_text: Optional[str] = None,
|
||||
summary_type: SummaryType = SummaryType.SHORT,
|
||||
) -> str:
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> Summary:
|
||||
"""
|
||||
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
|
||||
previous_summary_text: Optional previous summary to provide context
|
||||
summary_type: Type of summary to create ("short" or "long")
|
||||
max_tokens: Optional maximum number of tokens to generate. Supersedes summary_type.
|
||||
|
||||
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
|
||||
]
|
||||
)
|
||||
# Combine messages into a conversation format string
|
||||
conversation = format_messages(messages)
|
||||
|
||||
# Adjust system prompt based on summary type
|
||||
if summary_type == SummaryType.LONG:
|
||||
|
|
@ -182,9 +156,9 @@ but brief enough to be helpful.
|
|||
Return only the summary without any explanation or meta-commentary."""
|
||||
|
||||
# Include previous summary if available
|
||||
if previous_summary:
|
||||
if previous_summary_text:
|
||||
user_prompt = f"""Here is a previous summary of the conversation:
|
||||
{previous_summary}
|
||||
{previous_summary_text}
|
||||
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.
|
||||
|
|
@ -201,167 +175,205 @@ Provide a {"comprehensive" if summary_type == SummaryType.LONG else "concise"} s
|
|||
# Generate the summary
|
||||
llm_messages = [{"role": "user", "content": user_prompt}]
|
||||
|
||||
max_tokens = max_tokens or (1000 if summary_type == SummaryType.SHORT else 2000)
|
||||
|
||||
try:
|
||||
summary = await client.generate(
|
||||
summary_text = 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
|
||||
max_tokens=max_tokens,
|
||||
temperature=0.0,
|
||||
use_caching=True,
|
||||
)
|
||||
return summary
|
||||
logger.info("Successfully generated summary for session")
|
||||
return Summary(
|
||||
content=summary_text,
|
||||
message_count=len(messages),
|
||||
summary_type=summary_type.value,
|
||||
created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
message_id=messages[-1].id,
|
||||
# cheating: just use the max tokens instead of counting
|
||||
# tokens in response. in the future we can update client
|
||||
# to return tokens if that's doable.
|
||||
token_count=max_tokens,
|
||||
)
|
||||
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]}..."
|
||||
# Do not save this failed summary to the session metadata.
|
||||
return Summary(
|
||||
content=f"Conversation with {len(messages)} messages about {messages[-1].content[:30]}..."
|
||||
if messages
|
||||
else "No messages to summarize!",
|
||||
message_count=0,
|
||||
summary_type=summary_type.value,
|
||||
created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
message_id=messages[-1].id if messages else 0,
|
||||
token_count=50,
|
||||
)
|
||||
|
||||
|
||||
async def save_summary_metamessage(
|
||||
async def save_summary(
|
||||
db: AsyncSession,
|
||||
app_id: str,
|
||||
user_id: str,
|
||||
session_id: str,
|
||||
message_id: str,
|
||||
summary_content: str,
|
||||
message_count: int,
|
||||
summary_type: SummaryType = SummaryType.SHORT,
|
||||
) -> models.Metamessage:
|
||||
summary: Summary,
|
||||
workspace_name: str,
|
||||
session_name: str,
|
||||
) -> None:
|
||||
"""
|
||||
Save a summary as a metamessage.
|
||||
Save a summary as metadata on a session.
|
||||
|
||||
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
|
||||
summary: The summary to save
|
||||
workspace_name: Workspace name
|
||||
session_name: Session name
|
||||
|
||||
Returns:
|
||||
The created metamessage
|
||||
The updated session
|
||||
"""
|
||||
from src.exceptions import ResourceNotFoundException
|
||||
|
||||
# Get the label value from the enum
|
||||
label_value = summary_type.value
|
||||
label_value = summary["summary_type"]
|
||||
|
||||
# Create and save the metamessage
|
||||
metamessage = models.Metamessage(
|
||||
app_id=app_id,
|
||||
user_id=user_id,
|
||||
session_id=session_id,
|
||||
message_id=message_id,
|
||||
label=label_value,
|
||||
content=summary_content,
|
||||
h_metadata={"message_count": message_count, "summary_type": summary_type.name},
|
||||
)
|
||||
try:
|
||||
session = await crud.get_session(db, session_name, workspace_name)
|
||||
except ResourceNotFoundException:
|
||||
# If session doesn't exist, we can't save the summary
|
||||
logger.warning(
|
||||
f"Cannot save summary: session {session_name} not found in workspace {workspace_name}"
|
||||
)
|
||||
return
|
||||
|
||||
# Get existing summaries or create new dict
|
||||
existing_summaries = session.internal_metadata.get("summaries", {})
|
||||
existing_summaries[label_value] = summary
|
||||
|
||||
# Update the object metadata - create new dict to ensure SQLAlchemy detects the change
|
||||
updated_metadata = session.internal_metadata.copy()
|
||||
updated_metadata["summaries"] = existing_summaries
|
||||
session.internal_metadata = updated_metadata
|
||||
|
||||
db.add(metamessage)
|
||||
await db.commit()
|
||||
|
||||
logger.info(
|
||||
f"Saved {summary_type.name.lower()} summary metamessage for session {session_id} covering {message_count} messages"
|
||||
"Saved %s for session %s covering %s messages",
|
||||
summary["summary_type"],
|
||||
session_name,
|
||||
summary["message_count"],
|
||||
)
|
||||
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]]:
|
||||
db: AsyncSession,
|
||||
workspace_name: str,
|
||||
session_name: str,
|
||||
peer_name: str,
|
||||
cutoff: int | None = None,
|
||||
summary_type: SummaryType = SummaryType.SHORT,
|
||||
) -> str:
|
||||
"""
|
||||
Get a summarized version of the chat history by combining the latest summary
|
||||
with all messages since that summary.
|
||||
|
||||
Note: history is exclusive of the cutoff message.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
session_id: The session ID
|
||||
workspace_name: The workspace name
|
||||
session_name: The session name
|
||||
peer_name: The peer name
|
||||
cutoff: (Optional) message ID to cutoff at
|
||||
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
|
||||
A string formatted history text with summary and recent messages
|
||||
"""
|
||||
# Get messages since the latest summary and the summary itself
|
||||
messages, latest_summary = await get_messages_since_latest_summary(
|
||||
db, session_id, summary_type
|
||||
messages, latest_summary = await get_latest_summary_and_messages_since(
|
||||
db, workspace_name, session_name, peer_name, cutoff, 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}"
|
||||
return 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
|
||||
return messages_text
|
||||
|
||||
|
||||
async def get_messages_since_latest_summary(
|
||||
db: AsyncSession, session_id: str, summary_type: SummaryType = SummaryType.SHORT
|
||||
) -> tuple[list[models.Message], Optional[models.Metamessage]]:
|
||||
async def get_latest_summary_and_messages_since(
|
||||
db: AsyncSession,
|
||||
workspace_name: str,
|
||||
session_name: str,
|
||||
peer_name: str,
|
||||
cutoff: int | None = None,
|
||||
summary_type: SummaryType = SummaryType.SHORT,
|
||||
) -> tuple[list[models.Message], Optional[Summary]]:
|
||||
"""
|
||||
Get all messages since the latest summary for a session.
|
||||
Get all messages since the latest summary for a session or peer.
|
||||
|
||||
This is a convenience method that combines:
|
||||
1. Getting the latest summary for the session
|
||||
1. Getting the latest summary for the session or peer
|
||||
2. Getting all messages since that summary
|
||||
|
||||
Note that if the latest summary is not found, this will return all messages
|
||||
since the start of the session.
|
||||
|
||||
Note: history is exclusive of the cutoff message.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
session_id: The session ID
|
||||
workspace_name: The workspace name
|
||||
session_name: The session name
|
||||
peer_name: The peer name
|
||||
cutoff: (Optional) message ID to cutoff at
|
||||
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
|
||||
- The latest summary data, 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)
|
||||
# Get the latest summary
|
||||
summary = await get_summary(db, workspace_name, session_name, summary_type)
|
||||
|
||||
# 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
|
||||
if summary:
|
||||
messages = await crud.get_messages_id_range(
|
||||
db,
|
||||
workspace_name,
|
||||
session_name,
|
||||
peer_name,
|
||||
start_id=summary["message_id"],
|
||||
end_id=cutoff,
|
||||
)
|
||||
return messages, latest_summary
|
||||
return messages, summary
|
||||
else:
|
||||
messages = await get_messages_since_message(db, session_id)
|
||||
messages = await crud.get_messages_id_range(
|
||||
db, workspace_name, session_name, peer_name, end_id=cutoff
|
||||
)
|
||||
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]]:
|
||||
db: AsyncSession,
|
||||
workspace_name: str,
|
||||
session_name: str,
|
||||
peer_name: str,
|
||||
message_id: int,
|
||||
summary_type: SummaryType = SummaryType.SHORT,
|
||||
) -> tuple[bool, list[models.Message], Optional[Summary]]:
|
||||
"""
|
||||
Determine if a new summary should be created for this session.
|
||||
Determine if a new summary should be created for this object (peer or session).
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
session_id: The session ID
|
||||
workspace_name: The workspace name
|
||||
session_name: The session name
|
||||
summary_type: Type of summary to check for ("short" or "long")
|
||||
|
||||
Returns:
|
||||
|
|
@ -370,8 +382,13 @@ async def should_create_summary(
|
|||
- 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
|
||||
messages, latest_summary = await get_latest_summary_and_messages_since(
|
||||
db,
|
||||
workspace_name,
|
||||
session_name,
|
||||
peer_name,
|
||||
cutoff=message_id,
|
||||
summary_type=summary_type,
|
||||
)
|
||||
threshold = (
|
||||
MESSAGES_PER_SHORT_SUMMARY
|
||||
|
|
@ -379,15 +396,20 @@ async def should_create_summary(
|
|||
else MESSAGES_PER_LONG_SUMMARY
|
||||
)
|
||||
should_create = len(messages) >= threshold
|
||||
logger.debug(
|
||||
"Should create summary: %s, messages: %s, threshold: %s",
|
||||
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.
|
||||
Format a list of messages into a string by concatenating their content and
|
||||
prefixing each with the peer name.
|
||||
"""
|
||||
if len(messages) == 0:
|
||||
return ""
|
||||
return "\n".join(
|
||||
[f"{'user' if msg.is_user else 'assistant'}: {msg.content}" for msg in messages]
|
||||
)
|
||||
return "\n".join([f"{msg.peer_name}: {msg.content}" for msg in messages])
|
||||
|
|
|
|||
|
|
@ -204,16 +204,18 @@ def auth_client(client, request, monkeypatch):
|
|||
async def sample_data(db_session):
|
||||
"""Helper function to create test data"""
|
||||
# Create test app
|
||||
test_app = models.App(name=str(generate_nanoid()))
|
||||
db_session.add(test_app)
|
||||
test_workspace = models.Workspace(name=str(generate_nanoid()))
|
||||
db_session.add(test_workspace)
|
||||
await db_session.flush()
|
||||
|
||||
# Create test user
|
||||
test_user = models.User(name=str(generate_nanoid()), app_id=test_app.public_id)
|
||||
db_session.add(test_user)
|
||||
test_peer = models.Peer(
|
||||
name=str(generate_nanoid()), workspace_name=test_workspace.name
|
||||
)
|
||||
db_session.add(test_peer)
|
||||
await db_session.flush()
|
||||
|
||||
yield test_app, test_user
|
||||
yield test_workspace, test_peer
|
||||
|
||||
await db_session.rollback()
|
||||
|
||||
|
|
@ -254,3 +256,122 @@ def mock_openai_embeddings():
|
|||
mock_response.data = [MagicMock(embedding=[0.1] * 1536)]
|
||||
mock_create.return_value = mock_response
|
||||
yield mock_create
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_model_client(request):
|
||||
"""Mock ModelClient to avoid needing API keys during tests"""
|
||||
# Skip mocking for ModelClient unit tests
|
||||
if "test_model_client" in request.node.name or "test_model_client.py" in str(
|
||||
request.fspath
|
||||
):
|
||||
yield None
|
||||
return
|
||||
|
||||
# Create a mock instance
|
||||
mock_client_instance = MagicMock()
|
||||
mock_client_instance.generate = AsyncMock(return_value="Test summary content")
|
||||
mock_client_instance.stream = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("src.utils.history.ModelClient") as mock_history_client,
|
||||
patch("src.deriver.tom.single_prompt.ModelClient") as mock_single_prompt_client,
|
||||
patch("src.deriver.tom.long_term.ModelClient") as mock_long_term_client,
|
||||
patch("src.agent.ModelClient") as mock_agent_client,
|
||||
):
|
||||
# Make all class constructors return our mock instance
|
||||
mock_history_client.return_value = mock_client_instance
|
||||
mock_single_prompt_client.return_value = mock_client_instance
|
||||
mock_long_term_client.return_value = mock_client_instance
|
||||
mock_agent_client.return_value = mock_client_instance
|
||||
|
||||
yield {
|
||||
"history": mock_history_client,
|
||||
"single_prompt": mock_single_prompt_client,
|
||||
"long_term": mock_long_term_client,
|
||||
"agent": mock_agent_client,
|
||||
"instance": mock_client_instance,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_tracked_db(db_session):
|
||||
"""Mock tracked_db to use the test database session"""
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
@asynccontextmanager
|
||||
async def mock_tracked_db_context(operation_name=None):
|
||||
yield db_session
|
||||
|
||||
with patch("src.dependencies.tracked_db", mock_tracked_db_context):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_crud_collection_operations():
|
||||
"""Mock CRUD operations that try to commit to database during tests"""
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
||||
from src import models
|
||||
|
||||
async def mock_get_or_create_collection(
|
||||
db, workspace_name, peer_name, collection_name
|
||||
):
|
||||
# Create a mock collection object that doesn't require database commit
|
||||
mock_collection = models.Collection(
|
||||
name="honcho",
|
||||
workspace_name=workspace_name,
|
||||
peer_name=peer_name,
|
||||
)
|
||||
mock_collection.id = generate_nanoid()
|
||||
return mock_collection
|
||||
|
||||
with patch(
|
||||
"src.crud.get_or_create_collection",
|
||||
mock_get_or_create_collection,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_agent_api_calls(request):
|
||||
"""Mock API calls made by the agent during tests"""
|
||||
# Mock the agent-specific functions
|
||||
with (
|
||||
patch("src.agent.generate_semantic_queries") as mock_generate_queries,
|
||||
patch("src.deriver.tom.get_tom_inference") as mock_tom_inference,
|
||||
patch("src.agent.get_user_representation_long_term") as mock_user_rep,
|
||||
patch(
|
||||
"src.deriver.tom.embeddings.CollectionEmbeddingStore.get_relevant_facts"
|
||||
) as mock_get_facts,
|
||||
patch("src.agent.Dialectic.call") as mock_dialectic_call,
|
||||
patch("src.agent.Dialectic.stream") as mock_dialectic_stream,
|
||||
):
|
||||
# Mock semantic query generation
|
||||
mock_generate_queries.return_value = ["test query 1", "test query 2"]
|
||||
|
||||
# Mock ToM inference
|
||||
mock_tom_inference.return_value = (
|
||||
"<prediction>Test prediction about user mental state</prediction>"
|
||||
)
|
||||
|
||||
# Mock user representation generation
|
||||
mock_user_rep.return_value = (
|
||||
"<representation>Test user representation</representation>"
|
||||
)
|
||||
|
||||
# Mock embedding store facts retrieval
|
||||
mock_get_facts.return_value = ["fact 1", "fact 2", "fact 3"]
|
||||
|
||||
# Mock Dialectic API calls
|
||||
mock_dialectic_call.return_value = [{"text": "Test dialectic response"}]
|
||||
mock_dialectic_stream.return_value = AsyncMock()
|
||||
|
||||
yield {
|
||||
"generate_queries": mock_generate_queries,
|
||||
"tom_inference": mock_tom_inference,
|
||||
"user_rep": mock_user_rep,
|
||||
"get_facts": mock_get_facts,
|
||||
"dialectic_call": mock_dialectic_call,
|
||||
"dialectic_stream": mock_dialectic_stream,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,599 @@
|
|||
from typing import Optional
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
from sqlalchemy import select
|
||||
|
||||
from src import crud, models, schemas
|
||||
from src.models import QueueItem
|
||||
from src.routers.messages import enqueue
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestEnqueueFunction:
|
||||
"""Test suite for the enqueue function's internal logic"""
|
||||
|
||||
# Helper methods
|
||||
def create_sample_payload(
|
||||
self,
|
||||
workspace_name="test_workspace",
|
||||
session_name: Optional[str] = "test_session",
|
||||
peer_name="test_peer",
|
||||
count=1,
|
||||
):
|
||||
"""Create sample payload for testing"""
|
||||
return [
|
||||
{
|
||||
"workspace_name": workspace_name,
|
||||
"session_name": session_name,
|
||||
"message_id": f"msg_{i}",
|
||||
"content": f"Test message {i}",
|
||||
"metadata": {"test": f"value_{i}"},
|
||||
"peer_name": peer_name,
|
||||
}
|
||||
for i in range(count)
|
||||
]
|
||||
|
||||
async def count_queue_items(self, db_session):
|
||||
"""Helper to count queue items in database"""
|
||||
result = await db_session.execute(select(QueueItem))
|
||||
return len(result.scalars().all())
|
||||
|
||||
# Edge Cases & Input Validation Tests
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_payload_skips_enqueue(self, caplog):
|
||||
"""Test that empty payload list is handled gracefully"""
|
||||
with caplog.at_level("DEBUG"):
|
||||
await enqueue([])
|
||||
|
||||
assert "Empty payload list, skipping enqueue" in caplog.text
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_malformed_payload_logs_error(self, caplog):
|
||||
"""Test that malformed payload logs appropriate error"""
|
||||
malformed_payload = [{"incomplete": "data"}]
|
||||
|
||||
with caplog.at_level("ERROR"):
|
||||
await enqueue(malformed_payload) # Should not raise, but log error
|
||||
|
||||
assert "Failed to enqueue messages: 'workspace_name'" in caplog.text
|
||||
|
||||
# PEER MESSAGE TESTS
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_none_session_with_observe_me_true(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
"""Test peer-only processing when observe_me=True"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=None,
|
||||
peer_name=test_peer.name,
|
||||
count=1,
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create 1 queue item (representation)
|
||||
assert final_count - initial_count == 1
|
||||
|
||||
# Verify queue items have correct structure
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id.is_(None))
|
||||
)
|
||||
queue_items = result.scalars().all()
|
||||
|
||||
for item in queue_items:
|
||||
payload_data = item.payload
|
||||
assert payload_data["task_type"] == "representation"
|
||||
assert payload_data["sender_name"] == test_peer.name
|
||||
assert payload_data["target_name"] == test_peer.name
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_none_session_with_observe_me_false(
|
||||
self, mock_tracked_db, db_session, sample_data, caplog
|
||||
):
|
||||
"""Test peer-only processing when observe_me=False"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
test_peer.configuration = {"observe_me": False}
|
||||
await db_session.commit()
|
||||
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=None,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
with caplog.at_level("INFO"):
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should not create any queue items
|
||||
assert final_count == initial_count
|
||||
assert (
|
||||
f"Peer {test_peer.name} has observe_me=False, skipping enqueue"
|
||||
in caplog.text
|
||||
)
|
||||
|
||||
# SESSION MESSAGES
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_session_with_deriver_disabled(
|
||||
self, mock_tracked_db, db_session, sample_data, caplog
|
||||
):
|
||||
"""Test that deriver disabled sessions skip enqueue"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create session with deriver disabled
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name,
|
||||
name=str(generate_nanoid()),
|
||||
configuration={"deriver_disabled": True},
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# we expect 1 queue item for session level summarization
|
||||
assert final_count == initial_count + 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_session_normal_processing_single_peer(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
test_session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
schemas.SessionCreate(
|
||||
name=str(generate_nanoid()),
|
||||
peers={test_peer.name: schemas.SessionPeerConfig(observe_me=True)},
|
||||
),
|
||||
test_workspace.name,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create 1 queue items: 1 representation
|
||||
assert final_count - initial_count == 1
|
||||
|
||||
# Verify queue items
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id == test_session.id)
|
||||
)
|
||||
queue_items = result.scalars().all()
|
||||
|
||||
task_types = [item.payload["task_type"] for item in queue_items]
|
||||
assert "representation" in task_types
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_session_with_multiple_peers_none_observe_others(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
"""Test session processing with multiple peers where some observe others"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer1 = sample_data
|
||||
|
||||
# Create second peer
|
||||
test_peer2 = models.Peer(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_peer2)
|
||||
|
||||
test_session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
schemas.SessionCreate(
|
||||
name=str(generate_nanoid()),
|
||||
peers={
|
||||
test_peer1.name: schemas.SessionPeerConfig(),
|
||||
test_peer2.name: schemas.SessionPeerConfig(),
|
||||
},
|
||||
),
|
||||
test_workspace.name,
|
||||
)
|
||||
|
||||
await db_session.commit()
|
||||
|
||||
NUM_MESSAGES = 3
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer1.name,
|
||||
count=NUM_MESSAGES, # Message from peer1
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create NUM_MESSAGES * queue items:
|
||||
# 1 representation per sender
|
||||
assert final_count - initial_count == NUM_MESSAGES
|
||||
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id == test_session.id)
|
||||
)
|
||||
queue_items = result.scalars().all()
|
||||
|
||||
# Explicitly match up payloads by sender/target/task_type
|
||||
# For each message, we expect a representation and a summary queue item
|
||||
expected_payloads = []
|
||||
for _ in range(NUM_MESSAGES):
|
||||
expected_payloads.append(
|
||||
{
|
||||
"sender_name": test_peer1.name,
|
||||
"target_name": test_peer1.name,
|
||||
"task_type": "representation",
|
||||
}
|
||||
)
|
||||
actual_payloads = [
|
||||
{
|
||||
"sender_name": item.payload.get("sender_name"),
|
||||
"target_name": item.payload.get("target_name"),
|
||||
"task_type": item.payload.get("task_type"),
|
||||
}
|
||||
for item in queue_items
|
||||
]
|
||||
assert len(actual_payloads) == len(expected_payloads)
|
||||
|
||||
# Assert that all expected payloads are present in actual_payloads
|
||||
for expected in expected_payloads:
|
||||
assert expected in actual_payloads
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_session_with_multiple_peers_all_observe_others(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
"""Test session processing with multiple peers where some observe others"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer1 = sample_data
|
||||
|
||||
# Create second peer
|
||||
test_peer2 = models.Peer(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_peer2)
|
||||
|
||||
test_session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
schemas.SessionCreate(
|
||||
name=str(generate_nanoid()),
|
||||
peers={
|
||||
test_peer1.name: schemas.SessionPeerConfig(),
|
||||
test_peer2.name: schemas.SessionPeerConfig(observe_others=True),
|
||||
},
|
||||
),
|
||||
test_workspace.name,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
NUM_MESSAGES = 3
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer1.name,
|
||||
count=NUM_MESSAGES, # Message from peer1
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create NUM_MESSAGES * 2 queue items:
|
||||
# 1 representation for sender, 1 local representation for observer
|
||||
assert final_count - initial_count == NUM_MESSAGES * 2
|
||||
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id == test_session.id)
|
||||
)
|
||||
queue_items = result.scalars().all()
|
||||
|
||||
# Explicitly match up payloads by sender/target/task_type
|
||||
# For each message, we expect a representation and a summary queue item
|
||||
expected_payloads = []
|
||||
for _ in range(NUM_MESSAGES):
|
||||
expected_payloads.append(
|
||||
{
|
||||
"sender_name": test_peer1.name,
|
||||
"target_name": test_peer1.name,
|
||||
"task_type": "representation",
|
||||
}
|
||||
)
|
||||
expected_payloads.append(
|
||||
{
|
||||
"sender_name": test_peer1.name,
|
||||
"target_name": test_peer2.name,
|
||||
"task_type": "representation",
|
||||
}
|
||||
)
|
||||
actual_payloads = [
|
||||
{
|
||||
"sender_name": item.payload.get("sender_name"),
|
||||
"target_name": item.payload.get("target_name"),
|
||||
"task_type": item.payload.get("task_type"),
|
||||
}
|
||||
for item in queue_items
|
||||
]
|
||||
assert len(actual_payloads) == len(expected_payloads)
|
||||
|
||||
# Assert that all expected payloads are present in actual_payloads
|
||||
for expected in expected_payloads:
|
||||
assert expected in actual_payloads
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_session_with_multiple_peers_some_observe_others(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
"""Test session processing with multiple peers where some observe others"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer1 = sample_data
|
||||
|
||||
# Create second peer
|
||||
observing_peer = models.Peer(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(observing_peer)
|
||||
|
||||
# Create third peer
|
||||
unobserving_peer = models.Peer(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(unobserving_peer)
|
||||
|
||||
# Create session
|
||||
test_session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
schemas.SessionCreate(
|
||||
name=str(generate_nanoid()),
|
||||
peers={
|
||||
test_peer1.name: schemas.SessionPeerConfig(observe_me=True),
|
||||
observing_peer.name: schemas.SessionPeerConfig(observe_others=True),
|
||||
unobserving_peer.name: schemas.SessionPeerConfig(),
|
||||
},
|
||||
),
|
||||
test_workspace.name,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
NUM_MESSAGES = 3
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer1.name,
|
||||
count=NUM_MESSAGES, # Message from peer1
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create NUM_MESSAGES * 2 queue items:
|
||||
# 1 representation for sender, 1 local representation for 1 peer observer
|
||||
assert final_count - initial_count == NUM_MESSAGES * 2
|
||||
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id == test_session.id)
|
||||
)
|
||||
queue_items = result.scalars().all()
|
||||
|
||||
# Explicitly match up payloads by sender/target/task_type
|
||||
# For each message, we expect a representation and a summary queue item
|
||||
expected_payloads = []
|
||||
for _ in range(NUM_MESSAGES):
|
||||
expected_payloads.append(
|
||||
{
|
||||
"sender_name": test_peer1.name,
|
||||
"target_name": test_peer1.name,
|
||||
"task_type": "representation",
|
||||
}
|
||||
)
|
||||
expected_payloads.append(
|
||||
{
|
||||
"sender_name": test_peer1.name,
|
||||
"target_name": observing_peer.name,
|
||||
"task_type": "representation",
|
||||
}
|
||||
)
|
||||
actual_payloads = [
|
||||
{
|
||||
"sender_name": item.payload.get("sender_name"),
|
||||
"target_name": item.payload.get("target_name"),
|
||||
"task_type": item.payload.get("task_type"),
|
||||
}
|
||||
for item in queue_items
|
||||
]
|
||||
assert len(actual_payloads) == len(expected_payloads)
|
||||
|
||||
# Assert that all expected payloads are present in actual_payloads
|
||||
for expected in expected_payloads:
|
||||
assert expected in actual_payloads
|
||||
|
||||
assert unobserving_peer.name not in [
|
||||
item.payload.get("target_name") for item in queue_items
|
||||
]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_session_peer_config_overrides_peer_config(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
"""Test that session peer config overrides peer config"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Set peer configuration to observe_me=True
|
||||
test_peer.configuration = {"observe_me": True}
|
||||
|
||||
# Create session
|
||||
test_session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
schemas.SessionCreate(
|
||||
name=str(generate_nanoid()),
|
||||
peers={test_peer.name: schemas.SessionPeerConfig(observe_me=False)},
|
||||
),
|
||||
test_workspace.name,
|
||||
)
|
||||
|
||||
await db_session.commit()
|
||||
|
||||
payload = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create 0 queue items
|
||||
assert final_count - initial_count == 0
|
||||
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id == test_session.id)
|
||||
)
|
||||
queue_items = result.scalars().all()
|
||||
assert len(queue_items) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("src.routers.messages.tracked_db")
|
||||
async def test_multi_sender_scenario(
|
||||
self, mock_tracked_db, db_session, sample_data
|
||||
):
|
||||
"""Test complex scenario with multiple peers and mixed configurations"""
|
||||
mock_tracked_db.return_value.__aenter__.return_value = db_session
|
||||
|
||||
test_workspace, test_peer1 = sample_data
|
||||
|
||||
# Create additional peers
|
||||
additional_sender_peer = models.Peer(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
|
||||
observer_peer = models.Peer(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
|
||||
db_session.add_all([additional_sender_peer, observer_peer])
|
||||
|
||||
test_session = await crud.get_or_create_session(
|
||||
db_session,
|
||||
schemas.SessionCreate(
|
||||
name=str(generate_nanoid()),
|
||||
peers={
|
||||
test_peer1.name: schemas.SessionPeerConfig(),
|
||||
additional_sender_peer.name: schemas.SessionPeerConfig(),
|
||||
observer_peer.name: schemas.SessionPeerConfig(observe_others=True),
|
||||
},
|
||||
),
|
||||
test_workspace.name,
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
payload1 = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=test_peer1.name,
|
||||
)
|
||||
|
||||
payload2 = self.create_sample_payload(
|
||||
workspace_name=test_workspace.name,
|
||||
session_name=test_session.name,
|
||||
peer_name=additional_sender_peer.name,
|
||||
)
|
||||
|
||||
payload = payload1 + payload2
|
||||
|
||||
initial_count = await self.count_queue_items(db_session)
|
||||
await enqueue(payload)
|
||||
final_count = await self.count_queue_items(db_session)
|
||||
|
||||
# Should create 4 queue items:
|
||||
# 1 representation for test_peer1 (sender) and observer_peer (target)
|
||||
# 1 representation for additional_sender_peer (sender) and observer_peer (target)
|
||||
assert final_count - initial_count == 4
|
||||
|
||||
# Verify the correct representations were created
|
||||
result = await db_session.execute(
|
||||
select(QueueItem).where(QueueItem.session_id == test_session.id)
|
||||
)
|
||||
queue_items = result.all()
|
||||
|
||||
# Build expected payloads for this scenario
|
||||
expected_payloads = []
|
||||
# 2 messages: one from test_peer1, one from additional_sender_peer
|
||||
for sender in [test_peer1.name, additional_sender_peer.name]:
|
||||
# representation for sender (self)
|
||||
expected_payloads.append(
|
||||
{
|
||||
"task_type": "representation",
|
||||
"sender_name": sender,
|
||||
"target_name": sender,
|
||||
}
|
||||
)
|
||||
# representation for observer_peer (observe_others=True)
|
||||
expected_payloads.append(
|
||||
{
|
||||
"task_type": "representation",
|
||||
"sender_name": sender,
|
||||
"target_name": observer_peer.name,
|
||||
}
|
||||
)
|
||||
|
||||
# Extract actual payloads (task_type, sender_name, target_name) from queue_items
|
||||
actual_payloads = [
|
||||
{
|
||||
"task_type": item[0].payload.get("task_type"),
|
||||
"sender_name": item[0].payload.get("sender_name"),
|
||||
"target_name": item[0].payload.get("target_name"),
|
||||
}
|
||||
for item in queue_items
|
||||
]
|
||||
|
||||
assert len(actual_payloads) == len(expected_payloads)
|
||||
|
||||
# For each expected payload, assert it is present in actual_payloads
|
||||
for expected in expected_payloads:
|
||||
assert expected in actual_payloads
|
||||
|
|
@ -1,151 +0,0 @@
|
|||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_create_app(client):
|
||||
name = str(generate_nanoid())
|
||||
response = client.post(
|
||||
"/v1/apps", json={"name": name, "metadata": {"key": "value"}}
|
||||
)
|
||||
print(response)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
print("===================")
|
||||
print(data)
|
||||
print("===================")
|
||||
assert data["name"] == name
|
||||
assert data["metadata"] == {"key": "value"}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_create_app_no_metadata(client):
|
||||
name = str(generate_nanoid())
|
||||
response = client.post("/v1/apps", json={"name": name})
|
||||
print(response)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
print("===================")
|
||||
print(data)
|
||||
print("===================")
|
||||
assert data["name"] == name
|
||||
assert data["metadata"] == {}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_or_create_app(client):
|
||||
name = str(generate_nanoid())
|
||||
# Should return a ResourceNotFoundException with 404 status
|
||||
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
|
||||
data = response.json()
|
||||
assert data["name"] == name
|
||||
assert "id" in data
|
||||
|
||||
|
||||
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"]
|
||||
assert app1["metadata"] == app2["metadata"]
|
||||
|
||||
|
||||
def test_get_app_by_id(client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
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}")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == test_app.name
|
||||
assert data["id"] == str(test_app.public_id)
|
||||
|
||||
|
||||
def test_update_app(client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
new_name = str(generate_nanoid())
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}",
|
||||
json={"name": new_name, "metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == new_name
|
||||
assert data["metadata"] == {"new_key": "new_value"}
|
||||
|
||||
|
||||
def test_create_duplicate_app_name(client):
|
||||
# Create an app
|
||||
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()
|
||||
assert "detail" in data
|
||||
|
|
@ -1,185 +0,0 @@
|
|||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_create_collection(client, sample_data) -> None:
|
||||
test_app, test_user = sample_data
|
||||
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()
|
||||
assert data["name"] == "test_collection"
|
||||
assert data["metadata"] == {}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_collection_by_id(client, sample_data) -> None:
|
||||
test_app, test_user = sample_data
|
||||
# Make the collection
|
||||
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()
|
||||
# Get the collection
|
||||
response = client.get(
|
||||
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()
|
||||
assert data["name"] == "test_collection"
|
||||
assert data["metadata"] == {}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_collection_by_name(client, sample_data) -> None:
|
||||
test_app, test_user = sample_data
|
||||
# Make the collection
|
||||
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()
|
||||
# Get the collection
|
||||
response = client.get(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/name/test_collection"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "test_collection"
|
||||
assert data["metadata"] == {}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_collections(client, sample_data) -> None:
|
||||
test_app, test_user = sample_data
|
||||
# Make Sample Collections
|
||||
client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"test": "key"}},
|
||||
)
|
||||
client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"test": "key"}},
|
||||
)
|
||||
client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"test": "key2"}},
|
||||
)
|
||||
|
||||
# Get the Collections
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data["items"]) == 3
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/list",
|
||||
json={"filter": {"test": "key"}},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data["items"]) == 2
|
||||
assert data["items"][0]["metadata"] == {"test": "key"}
|
||||
assert data["items"][1]["metadata"] == {"test": "key"}
|
||||
|
||||
|
||||
def test_update_collection(client, sample_data) -> None:
|
||||
test_app, test_user = sample_data
|
||||
# Make the collection
|
||||
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()
|
||||
# Update the collection
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}",
|
||||
json={"name": "test_collection_updated", "metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == "test_collection_updated"
|
||||
assert data["metadata"] == {"new_key": "new_value"}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_delete_collection(client, sample_data) -> None:
|
||||
test_app, test_user = sample_data
|
||||
# Make the collection
|
||||
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()
|
||||
# Delete the collection
|
||||
response = client.delete(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
response = client.get(
|
||||
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()
|
||||
|
|
@ -1,261 +0,0 @@
|
|||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_update_document_validation_error(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
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
|
||||
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",
|
||||
json={"content": "test_text", "metadata": {}},
|
||||
)
|
||||
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()
|
||||
assert "detail" in data
|
||||
|
||||
|
||||
def test_create_document(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
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()
|
||||
# Create a document
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}/documents",
|
||||
json={"content": "test_text", "metadata": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["content"] == "test_text"
|
||||
assert data["metadata"] == {}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_document(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
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
|
||||
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",
|
||||
json={"content": "test_text", "metadata": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
document = response.json()
|
||||
# Get the document
|
||||
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
|
||||
data = response.json()
|
||||
assert data["content"] == "test_text"
|
||||
assert data["metadata"] == {}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_documents(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
|
||||
json={"name": str(generate_nanoid()), "metadata": {}},
|
||||
)
|
||||
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",
|
||||
json={"content": "test_text", "metadata": {"test": "key"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents",
|
||||
json={"content": "test_text", "metadata": {"test": "key"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents",
|
||||
json={"content": "test_text", "metadata": {"test": "key2"}},
|
||||
)
|
||||
# Get the documents
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data["items"]) == 3
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/list",
|
||||
json={"filter": {"test": "key"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data["items"]) == 2
|
||||
assert data["items"][0]["metadata"]["test"] == "key"
|
||||
assert data["items"][1]["metadata"]["test"] == "key"
|
||||
|
||||
|
||||
def test_query_documents(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
|
||||
json={"name": str(generate_nanoid()), "metadata": {}},
|
||||
)
|
||||
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",
|
||||
json={"content": "test_text", "metadata": {"test": "key"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents",
|
||||
json={"content": "test_text", "metadata": {"test": "key"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents",
|
||||
json={"content": "test_text", "metadata": {"test": "key2"}},
|
||||
)
|
||||
# Get the documents
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/query",
|
||||
json={"query": "test"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
print("=====================")
|
||||
print(data)
|
||||
print("=====================")
|
||||
assert len(data) == 3
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{collection['id']}/documents/query",
|
||||
json={"query": "test", "filter": {"test": "key"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 2
|
||||
assert data[0]["metadata"]["test"] == "key"
|
||||
assert data[1]["metadata"]["test"] == "key"
|
||||
|
||||
|
||||
def test_update_document(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
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()
|
||||
# Create a document
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}/documents",
|
||||
json={"content": "test_text", "metadata": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
# Update the document
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections/{data['id']}/documents/{data['id']}",
|
||||
json={"content": "test_text_updated", "metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
|
||||
|
||||
def test_delete_document(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a collection
|
||||
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
|
||||
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",
|
||||
json={"content": "test_text", "metadata": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
document = response.json()
|
||||
# Delete the document
|
||||
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 == 200
|
||||
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 == 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"
|
||||
|
|
@ -12,20 +12,22 @@ def test_create_key_no_params(auth_client):
|
|||
|
||||
def test_create_key_with_params(auth_client, sample_data):
|
||||
"""Test creating a key with specific parameters"""
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = 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})
|
||||
response = auth_client.post(
|
||||
"/v1/keys", params={"workspace_id": test_workspace.name}
|
||||
)
|
||||
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},
|
||||
params={"workspace_id": test_workspace.name, "peer_id": test_peer.name},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert "key" in response.json()
|
||||
|
|
@ -34,8 +36,8 @@ def test_create_key_with_params(auth_client, sample_data):
|
|||
response = auth_client.post(
|
||||
"/v1/keys",
|
||||
params={
|
||||
"app_id": test_app.public_id,
|
||||
"user_id": test_user.public_id,
|
||||
"workspace_id": test_workspace.name,
|
||||
"peer_id": test_peer.name,
|
||||
"session_id": "test-session",
|
||||
"collection_id": "test-collection",
|
||||
},
|
||||
|
|
@ -56,8 +58,10 @@ def test_create_key_with_expires_at(auth_client, sample_data):
|
|||
else:
|
||||
assert response.status_code == 401
|
||||
|
||||
test_app, _ = sample_data
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# assert that the key is expired
|
||||
response = auth_client.post("/v1/keys", params={"app_id": test_app.public_id})
|
||||
response = auth_client.post(
|
||||
"/v1/keys", params={"workspace_id": test_workspace.name}
|
||||
)
|
||||
assert response.status_code == 401
|
||||
|
|
|
|||
|
|
@ -1,47 +1,170 @@
|
|||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
||||
from src import models # Import your SQLAlchemy models
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_message(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(user_id=test_user.public_id, app_id=test_app.public_id)
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={
|
||||
"content": "Test message",
|
||||
"is_user": True,
|
||||
"metadata": {"message_key": "message_value"},
|
||||
"messages": [
|
||||
{
|
||||
"content": "Test message",
|
||||
"peer_id": test_peer.name,
|
||||
"metadata": {"message_key": "message_value"},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["content"] == "Test message"
|
||||
assert data["is_user"] is True
|
||||
assert data["metadata"] == {"message_key": "message_value"}
|
||||
assert "id" in data
|
||||
assert len(data) == 1
|
||||
message = data[0]
|
||||
assert message["content"] == "Test message"
|
||||
assert message["peer_id"] == test_peer.name
|
||||
assert message["session_id"] == test_session.name
|
||||
assert message["metadata"] == {"message_key": "message_value"}
|
||||
assert "id" in message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages_with_metadata(client, db_session, sample_data):
|
||||
"""Test batch message creation with metadata for each message"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{
|
||||
"content": "Message 1",
|
||||
"peer_id": test_peer.name,
|
||||
"metadata": {"type": "question", "priority": "high"},
|
||||
},
|
||||
{
|
||||
"content": "Message 2",
|
||||
"peer_id": test_peer.name,
|
||||
"metadata": {"type": "answer", "priority": "low"},
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 2
|
||||
|
||||
# Check first message
|
||||
assert data[0]["content"] == "Message 1"
|
||||
assert data[0]["metadata"] == {"type": "question", "priority": "high"}
|
||||
|
||||
# Check second message
|
||||
assert data[1]["content"] == "Message 2"
|
||||
assert data[1]["metadata"] == {"type": "answer", "priority": "low"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages_without_metadata(client, db_session, sample_data):
|
||||
"""Test batch message creation without metadata (should default to empty dict)"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{
|
||||
"content": "Message without metadata",
|
||||
"peer_id": test_peer.name,
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 1
|
||||
assert data[0]["content"] == "Message without metadata"
|
||||
assert data[0]["metadata"] == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages_with_null_metadata(
|
||||
client, db_session, sample_data
|
||||
):
|
||||
"""Test batch message creation with null metadata"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{
|
||||
"content": "Message with null metadata",
|
||||
"peer_id": test_peer.name,
|
||||
"metadata": None,
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 1
|
||||
assert data[0]["content"] == "Message with null metadata"
|
||||
assert data[0]["metadata"] == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_messages(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(user_id=test_user.public_id, app_id=test_app.public_id)
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id, content="Test message", is_user=True, app_id=test_app.public_id, user_id=test_user.public_id
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/list",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
|
@ -49,31 +172,180 @@ async def test_get_messages(client, db_session, sample_data):
|
|||
assert "items" in data
|
||||
assert len(data["items"]) > 0
|
||||
assert data["items"][0]["content"] == "Test message"
|
||||
assert data["items"][0]["is_user"] is True
|
||||
assert data["items"][0]["peer_id"] == test_peer.name
|
||||
assert data["items"][0]["session_id"] == test_session.name
|
||||
assert data["items"][0]["metadata"] == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_filtered_messages(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a test session and message
|
||||
test_session = models.Session(user_id=test_user.public_id, app_id=test_app.public_id)
|
||||
async def test_get_messages_with_reverse(client, db_session, sample_data):
|
||||
"""Test getting messages with reverse parameter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
# Create multiple messages to test ordering
|
||||
test_message1 = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="First message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
test_message2 = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Second message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message1)
|
||||
db_session.add(test_message2)
|
||||
await db_session.commit()
|
||||
|
||||
# Test normal order
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
normal_data = response.json()
|
||||
|
||||
# Test reversed order
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list?reverse=true",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
reversed_data = response.json()
|
||||
|
||||
# Both should have items
|
||||
assert len(normal_data["items"]) >= 2
|
||||
assert len(reversed_data["items"]) >= 2
|
||||
|
||||
# Order should be different (first item in normal should be last in reversed)
|
||||
if len(normal_data["items"]) > 1 and len(reversed_data["items"]) > 1:
|
||||
assert normal_data["items"][0]["id"] != reversed_data["items"][0]["id"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_messages_with_empty_filter(client, db_session, sample_data):
|
||||
"""Test getting messages with empty filter object"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list",
|
||||
json={"filter": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_messages_with_null_filter(client, db_session, sample_data):
|
||||
"""Test getting messages with null filter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list",
|
||||
json={"filter": None},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_messages_no_body(client, db_session, sample_data):
|
||||
"""Test getting messages without any request body"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_filtered_messages(client, db_session, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and messages
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"key": "value"},
|
||||
)
|
||||
test_message2 = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
session_name=test_session.name,
|
||||
content="Test message 2",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"key": "value2"},
|
||||
)
|
||||
db_session.add(test_message)
|
||||
|
|
@ -81,34 +353,93 @@ async def test_get_filtered_messages(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}/messages/list",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list",
|
||||
json={"filter": {"key": "value2"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
print(data)
|
||||
assert "items" in data
|
||||
assert len(data["items"]) > 0
|
||||
assert data["items"][0]["content"] == "Test message"
|
||||
assert data["items"][0]["is_user"] is True
|
||||
assert data["items"][0]["content"] == "Test message 2"
|
||||
assert data["items"][0]["peer_id"] == test_peer.name
|
||||
assert data["items"][0]["session_id"] == test_session.name
|
||||
assert data["items"][0]["metadata"] == {"key": "value2"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_message(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a test session and message
|
||||
test_session = models.Session(user_id=test_user.public_id, app_id=test_app.public_id)
|
||||
async def test_get_filtered_messages_with_complex_filter(
|
||||
client, db_session, sample_data
|
||||
):
|
||||
"""Test getting messages with complex metadata filter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and messages
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message1 = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Message 1",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"type": "question", "priority": "high", "category": "technical"},
|
||||
)
|
||||
test_message2 = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Message 2",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"type": "answer", "priority": "high", "category": "technical"},
|
||||
)
|
||||
test_message3 = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Message 3",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"type": "question", "priority": "low", "category": "general"},
|
||||
)
|
||||
db_session.add(test_message1)
|
||||
db_session.add(test_message2)
|
||||
db_session.add(test_message3)
|
||||
await db_session.commit()
|
||||
|
||||
# Filter by multiple criteria
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/list",
|
||||
json={"filter": {"priority": "high", "category": "technical"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
# Should return messages 1 and 2 (both have high priority and technical category)
|
||||
assert len(data["items"]) >= 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_message(client, db_session, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id, content="Test message", is_user=True, app_id=test_app.public_id, user_id=test_user.public_id
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/{test_message.public_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{test_message.public_id}",
|
||||
json={"metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
|
@ -117,101 +448,271 @@ async def test_update_message(client, db_session, sample_data):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_message_empty_metadata(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
async def test_update_message_with_complex_metadata(client, db_session, sample_data):
|
||||
"""Test updating message with complex metadata structure"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(user_id=test_user.public_id, app_id=test_app.public_id)
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id, content="Test message", is_user=True, app_id=test_app.public_id, user_id=test_user.public_id
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
complex_metadata = {
|
||||
"tags": ["important", "follow-up"],
|
||||
"score": 8.5,
|
||||
"nested": {"category": "technical", "subcategory": "api"},
|
||||
"processed": True,
|
||||
}
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{test_message.public_id}",
|
||||
json={"metadata": complex_metadata},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == complex_metadata
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_message_empty_metadata(client, db_session, sample_data):
|
||||
# note that this should not change the metadata of the message.
|
||||
# this test is to ensure that the metadata is not changed when it is set to None.
|
||||
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"test_key": "test_value"},
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/{test_message.public_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{test_message.public_id}",
|
||||
json={"metadata": None},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
assert response.status_code == 200
|
||||
|
||||
# now ensure that the metadata is not changed
|
||||
response = client.get(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{test_message.public_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
print(data)
|
||||
# assert data["detail"] == "Message metadata cannot be empty"
|
||||
assert data["metadata"] == {"test_key": "test_value"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages(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, app_id=test_app.public_id)
|
||||
async def test_update_message_with_empty_dict_metadata(client, db_session, sample_data):
|
||||
"""Test updating message with empty dictionary metadata"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
# Create batch of test messages
|
||||
test_messages = {
|
||||
"messages": [
|
||||
{
|
||||
"content": f"Test message {i}",
|
||||
"is_user": i % 2 == 0, # Alternating user/non-user messages
|
||||
"metadata": {"batch_index": i},
|
||||
}
|
||||
for i in range(3)
|
||||
]
|
||||
}
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/batch",
|
||||
json=test_messages,
|
||||
test_message = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
h_metadata={"old_key": "old_value"},
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{test_message.public_id}",
|
||||
json={"metadata": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == {}
|
||||
|
||||
# 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}"
|
||||
assert message["is_user"] == (i % 2 == 0)
|
||||
assert message["metadata"] == {"batch_index": i}
|
||||
assert "id" in message
|
||||
assert message["session_id"] == test_session.public_id
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_single_message(client, db_session, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Verify messages were actually saved to the database
|
||||
# Create a test session and message
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
test_message = models.Message(
|
||||
session_name=test_session.name,
|
||||
content="Test message",
|
||||
workspace_name=test_workspace.name,
|
||||
peer_name=test_peer.name,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.get(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{test_message.public_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["content"] == "Test message"
|
||||
assert data["peer_id"] == test_peer.name
|
||||
assert data["session_id"] == test_session.name
|
||||
assert data["workspace_id"] == test_workspace.name
|
||||
assert data["id"] == test_message.public_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_nonexistent_message(client, db_session, sample_data):
|
||||
"""Test getting a message that doesn't exist"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
nonexistent_message_id = str(generate_nanoid())
|
||||
response = client.get(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{nonexistent_message_id}"
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_nonexistent_message(client, db_session, sample_data):
|
||||
"""Test updating a message that doesn't exist"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
nonexistent_message_id = str(generate_nanoid())
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages/{nonexistent_message_id}",
|
||||
json={"metadata": {"key": "value"}},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_messages_for_nonexistent_session(client, sample_data):
|
||||
"""Test creating messages for a session that doesn't exist - should create the session"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
nonexistent_session_id = str(generate_nanoid())
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/list",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{nonexistent_session_id}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{
|
||||
"content": "Test message",
|
||||
"peer_id": test_peer.name,
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
# Should create the session and return 200 with the created message
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 1
|
||||
assert data[0]["content"] == "Test message"
|
||||
assert data[0]["workspace_id"] == test_workspace.name
|
||||
assert data[0]["session_id"] == nonexistent_session_id
|
||||
assert data[0]["peer_id"] == test_peer.name
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_messages_for_nonexistent_session(client, sample_data):
|
||||
"""Test getting messages for a session that doesn't exist - should return empty list"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
nonexistent_session_id = str(generate_nanoid())
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{nonexistent_session_id}/messages/list",
|
||||
json={},
|
||||
)
|
||||
# Should return 200 with empty results (session doesn't exist = no messages)
|
||||
assert response.status_code == 200
|
||||
saved_messages = response.json()["items"]
|
||||
assert len(saved_messages) == 3
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages_limit(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_session = models.Session(user_id=test_user.public_id, app_id=test_app.public_id)
|
||||
async def test_create_empty_batch_messages(client, db_session, sample_data):
|
||||
"""Test creating an empty batch of messages"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
# Create batch with more than 100 messages
|
||||
test_messages = {
|
||||
"messages": [
|
||||
{
|
||||
"content": f"Test message {i}",
|
||||
"is_user": i % 2 == 0,
|
||||
"metadata": {"batch_index": i},
|
||||
}
|
||||
for i in range(101) # 101 messages
|
||||
]
|
||||
}
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={"messages": []},
|
||||
)
|
||||
# Should return 422 for validation error (empty list not allowed)
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_batch_messages_max_limit(client, db_session, sample_data):
|
||||
"""Test creating batch messages at the maximum limit (100 messages)"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create a test session
|
||||
test_session = models.Session(
|
||||
workspace_name=test_workspace.name, name=str(generate_nanoid())
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
|
||||
# Create exactly 100 messages (the maximum allowed)
|
||||
messages = [
|
||||
{"content": f"Message {i}", "peer_id": test_peer.name, "metadata": {"index": i}}
|
||||
for i in range(100)
|
||||
]
|
||||
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{test_session.public_id}/messages/batch",
|
||||
json=test_messages,
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{test_session.name}/messages",
|
||||
json={"messages": messages},
|
||||
)
|
||||
|
||||
assert response.status_code == 422 # Validation error
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "messages" in data["detail"][0]["loc"] # Error should mention messages field
|
||||
assert len(data) == 100
|
||||
assert data[0]["content"] == "Message 0"
|
||||
assert data[99]["content"] == "Message 99"
|
||||
|
|
|
|||
|
|
@ -1,440 +0,0 @@
|
|||
import pytest
|
||||
|
||||
from src import models # Import your SQLAlchemy models
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_metamessage(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, app_id=test_app.public_id
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
response = client.post(
|
||||
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": {},
|
||||
"label": "test_type",
|
||||
},
|
||||
)
|
||||
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"] == {}
|
||||
assert data["label"] == "test_type"
|
||||
assert data["metamessage_type"] == "test_type"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_metamessage(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, app_id=test_app.public_id
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
test_metamessage = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session.public_id,
|
||||
message_id=test_message.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="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}/metamessages/{test_metamessage.public_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["user_id"] == str(test_user.public_id)
|
||||
assert data["app_id"] == str(test_app.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"] == {}
|
||||
assert data["label"] == "test_type"
|
||||
assert data["metamessage_type"] == "test_type"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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, app_id=test_app.public_id
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
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,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session.public_id,
|
||||
message_id=test_message.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type",
|
||||
)
|
||||
test_metamessage_2 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session.public_id,
|
||||
message_id=test_message.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type",
|
||||
)
|
||||
test_metamessage_3 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session.public_id,
|
||||
message_id=test_message.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type",
|
||||
)
|
||||
test_metamessage_4 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session.public_id,
|
||||
message_id=test_message.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type_2",
|
||||
)
|
||||
db_session.add(test_metamessage_1)
|
||||
db_session.add(test_metamessage_2)
|
||||
db_session.add(test_metamessage_3)
|
||||
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}/metamessages/list",
|
||||
json={
|
||||
"session_id": str(test_session.public_id),
|
||||
"label": "test_type",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) == 3
|
||||
assert data["items"][0]["content"] == "Test Metamessage"
|
||||
assert data["items"][0]["label"] == "test_type"
|
||||
assert data["items"][0]["metamessage_type"] == "test_type"
|
||||
assert data["items"][0]["session_id"] == str(test_session.public_id)
|
||||
assert data["items"][0]["metadata"] == {}
|
||||
assert data["items"][0]["app_id"] == str(test_app.public_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_metamessage_by_user(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create 3 test sessions
|
||||
test_session_1 = models.Session(
|
||||
user_id=test_user.public_id, app_id=test_app.public_id
|
||||
)
|
||||
test_session_2 = models.Session(
|
||||
user_id=test_user.public_id, app_id=test_app.public_id
|
||||
)
|
||||
test_session_3 = models.Session(
|
||||
user_id=test_user.public_id, app_id=test_app.public_id
|
||||
)
|
||||
db_session.add(test_session_1)
|
||||
db_session.add(test_session_2)
|
||||
db_session.add(test_session_3)
|
||||
await db_session.commit()
|
||||
|
||||
# Create a message in each session
|
||||
test_message_1 = models.Message(
|
||||
session_id=test_session_1.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
test_message_2 = models.Message(
|
||||
session_id=test_session_2.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
test_message_3 = models.Message(
|
||||
session_id=test_session_3.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
db_session.add(test_message_1)
|
||||
db_session.add(test_message_2)
|
||||
db_session.add(test_message_3)
|
||||
await db_session.commit()
|
||||
|
||||
# Create metamessages across different sessions
|
||||
test_metamessage_1 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session_1.public_id,
|
||||
message_id=test_message_1.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type",
|
||||
)
|
||||
test_metamessage_2 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session_2.public_id,
|
||||
message_id=test_message_2.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type",
|
||||
)
|
||||
test_metamessage_3 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session_3.public_id,
|
||||
message_id=test_message_3.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type",
|
||||
)
|
||||
test_metamessage_4 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session_3.public_id,
|
||||
message_id=test_message_3.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="test_type_2",
|
||||
)
|
||||
# Create a user-level metamessage (no session/message)
|
||||
test_metamessage_5 = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
content="User level metamessage",
|
||||
h_metadata={},
|
||||
label="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={"label": "test_type"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
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]["label"] == "test_type"
|
||||
assert data["items"][0]["metamessage_type"] == "test_type"
|
||||
assert data["items"][0]["user_id"] == str(test_user.public_id)
|
||||
assert data["items"][0]["app_id"] == str(test_app.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"},
|
||||
"label": "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"
|
||||
assert data["app_id"] == str(test_app.public_id)
|
||||
assert data["label"] == "user_insight"
|
||||
assert data["metamessage_type"] == "user_insight"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_metamessage(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, app_id=test_app.public_id
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
content="Test message",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
test_metamessage = models.Metamessage(
|
||||
user_id=test_user.public_id,
|
||||
app_id=test_app.public_id,
|
||||
session_id=test_session.public_id,
|
||||
message_id=test_message.public_id,
|
||||
content="Test Metamessage",
|
||||
h_metadata={},
|
||||
label="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}/metamessages/{test_metamessage.public_id}",
|
||||
json={
|
||||
"metadata": {"new_key": "new_value"},
|
||||
"label": "updated_type",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == {"new_key": "new_value"}
|
||||
assert data["label"] == "updated_type"
|
||||
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)
|
||||
assert data["app_id"] == str(test_app.public_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_metamessage_with_label_and_alias(client, db_session, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
# Create a common session and message for both test cases
|
||||
test_session = models.Session(
|
||||
user_id=test_user.public_id, app_id=test_app.public_id
|
||||
)
|
||||
db_session.add(test_session)
|
||||
await db_session.commit()
|
||||
test_message = models.Message(
|
||||
session_id=test_session.public_id,
|
||||
content="Shared message content",
|
||||
is_user=True,
|
||||
app_id=test_app.public_id,
|
||||
user_id=test_user.public_id,
|
||||
)
|
||||
db_session.add(test_message)
|
||||
await db_session.commit()
|
||||
|
||||
common_payload_parts = {
|
||||
"session_id": str(test_session.public_id),
|
||||
"message_id": str(test_message.public_id),
|
||||
"metadata": {"source": "input_test"},
|
||||
}
|
||||
test_value_for_type = "input_consistency_test_type"
|
||||
|
||||
# 1. Create metamessage using "label"
|
||||
response_with_label = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
|
||||
json={
|
||||
**common_payload_parts,
|
||||
"content": "Content created with label",
|
||||
"label": test_value_for_type,
|
||||
},
|
||||
)
|
||||
assert response_with_label.status_code == 200
|
||||
data_from_label_input = response_with_label.json()
|
||||
|
||||
# 2. Create metamessage using "metamessage_type" (alias)
|
||||
response_with_alias = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
|
||||
json={
|
||||
**common_payload_parts,
|
||||
"content": "Content created with alias",
|
||||
"metamessage_type": test_value_for_type, # Using alias
|
||||
},
|
||||
)
|
||||
assert response_with_alias.status_code == 200
|
||||
data_from_alias_input = response_with_alias.json()
|
||||
|
||||
# Assertions for the first response (created with "label")
|
||||
assert data_from_label_input["content"] == "Content created with label"
|
||||
assert data_from_label_input["label"] == test_value_for_type
|
||||
assert data_from_label_input["metamessage_type"] == test_value_for_type
|
||||
assert data_from_label_input["metadata"] == common_payload_parts["metadata"]
|
||||
assert data_from_label_input["metamessage_type"] == test_value_for_type
|
||||
assert data_from_label_input["app_id"] == str(test_app.public_id)
|
||||
assert data_from_label_input["user_id"] == str(test_user.public_id)
|
||||
assert data_from_label_input["session_id"] == common_payload_parts["session_id"]
|
||||
assert data_from_label_input["message_id"] == common_payload_parts["message_id"]
|
||||
|
||||
# Assertions for the second response (created with "metamessage_type")
|
||||
assert data_from_alias_input["content"] == "Content created with alias"
|
||||
assert (
|
||||
data_from_alias_input["label"] == test_value_for_type
|
||||
) # Output should still be "label"
|
||||
assert data_from_alias_input["metamessage_type"] == test_value_for_type
|
||||
assert data_from_alias_input["metadata"] == common_payload_parts["metadata"]
|
||||
assert data_from_alias_input["metamessage_type"] == test_value_for_type
|
||||
assert data_from_alias_input["user_id"] == str(test_user.public_id)
|
||||
assert data_from_alias_input["session_id"] == common_payload_parts["session_id"]
|
||||
assert data_from_alias_input["message_id"] == common_payload_parts["message_id"]
|
||||
assert data_from_alias_input["app_id"] == str(test_app.public_id)
|
||||
# Key assertion: The output for the type/label field is consistent ("label") and has the correct value
|
||||
assert data_from_label_input["label"] == data_from_alias_input["label"]
|
||||
assert (
|
||||
data_from_label_input["metamessage_type"]
|
||||
== data_from_alias_input["metamessage_type"]
|
||||
)
|
||||
assert data_from_label_input["label"] == data_from_label_input["metamessage_type"]
|
||||
assert data_from_alias_input["label"] == data_from_alias_input["metamessage_type"]
|
||||
|
|
@ -0,0 +1,525 @@
|
|||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_get_or_create_peer(client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
name = str(generate_nanoid())
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": name, "metadata": {"peer_key": "peer_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
assert data["metadata"] == {"peer_key": "peer_value"}
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_or_create_peer_with_configuration(client, sample_data):
|
||||
"""Test peer creation with configuration parameter"""
|
||||
test_workspace, _ = sample_data
|
||||
name = str(generate_nanoid())
|
||||
configuration = {"experimental": True, "beta": False}
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": name, "configuration": configuration},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_get_or_create_peer_with_all_optional_params(client, sample_data):
|
||||
"""Test peer creation with all optional parameters"""
|
||||
test_workspace, _ = sample_data
|
||||
name = str(generate_nanoid())
|
||||
metadata = {"key": "value", "number": 42}
|
||||
configuration = {"feature1": True, "feature2": False}
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": name, "metadata": metadata, "configuration": configuration},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
assert data["metadata"] == metadata
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_get_or_create_existing_peer(client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
name = str(generate_nanoid())
|
||||
|
||||
# Create the peer
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": name, "metadata": {"peer_key": "peer_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
peer1 = response.json()
|
||||
|
||||
# Try to create the same peer again - should return existing peer
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": name, "metadata": {"peer_key": "peer_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
peer2 = response.json()
|
||||
|
||||
# Both should be the same peer
|
||||
assert peer1["id"] == peer2["id"]
|
||||
assert peer1["metadata"] == peer2["metadata"]
|
||||
|
||||
|
||||
def test_get_peers(client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# Create a few peers with metadata
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"peer_key": "peer_value"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"peer_key": "peer_value"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"peer_key": "peer_value2"}},
|
||||
)
|
||||
|
||||
# Get all peers
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) > 0
|
||||
|
||||
# Get peers with filter
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/list",
|
||||
json={"filter": {"peer_key": "peer_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) >= 2
|
||||
assert data["items"][0]["metadata"]["peer_key"] == "peer_value"
|
||||
|
||||
|
||||
def test_get_peers_with_empty_filter(client, sample_data):
|
||||
"""Test peer listing with empty filter object"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/list", json={"filter": {}}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_get_peers_with_null_filter(client, sample_data):
|
||||
"""Test peer listing with null filter"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/list", json={"filter": None}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_update_peer(client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == {"new_key": "new_value"}
|
||||
|
||||
|
||||
def test_update_peer_with_configuration(client, sample_data):
|
||||
"""Test peer update with configuration parameter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
configuration = {"new_feature": True, "legacy_feature": False}
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"configuration": configuration},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_update_peer_with_all_optional_params(client, sample_data):
|
||||
"""Test peer update with both metadata and configuration"""
|
||||
test_workspace, test_peer = sample_data
|
||||
metadata = {"updated_key": "updated_value", "count": 100}
|
||||
configuration = {"experimental": True, "beta": True}
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"metadata": metadata, "configuration": configuration},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == metadata
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_update_peer_with_null_metadata(client, sample_data):
|
||||
"""Test peer update with null metadata"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# First set some metadata
|
||||
client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"metadata": {"temp": "value"}},
|
||||
)
|
||||
|
||||
# Then clear it with null
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"metadata": None},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "metadata" in data
|
||||
|
||||
|
||||
def test_update_peer_with_null_configuration(client, sample_data):
|
||||
"""Test peer update with null configuration"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"configuration": None},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "configuration" in data
|
||||
|
||||
|
||||
def test_get_sessions_for_peer_no_sessions(client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Get sessions for the peer
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/sessions",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
|
||||
|
||||
def test_get_sessions_for_peer(client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create session for the peer
|
||||
session_name = str(generate_nanoid())
|
||||
create_response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_name, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
assert create_response.status_code == 200
|
||||
created_session = create_response.json()
|
||||
assert created_session["id"] == session_name
|
||||
|
||||
# Now get sessions for the peer and validate the session is returned
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/sessions",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
# Check that the created session is in the returned items
|
||||
session_ids = [item["id"] for item in data["items"]]
|
||||
assert session_name in session_ids
|
||||
assert len(data["items"]) == 1
|
||||
|
||||
|
||||
def test_get_sessions_for_peer_with_is_active_filter(client, sample_data):
|
||||
"""Test getting sessions for peer with is_active parameter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create and then delete a session to have inactive session
|
||||
session_name = str(generate_nanoid())
|
||||
client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_name, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
client.delete(f"/v1/workspaces/{test_workspace.name}/sessions/{session_name}")
|
||||
|
||||
# Test getting inactive sessions
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/sessions",
|
||||
json={"is_active": False},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
# Should find at least our deleted session
|
||||
inactive_sessions = [s for s in data["items"] if not s["is_active"]]
|
||||
assert len(inactive_sessions) > 0
|
||||
|
||||
|
||||
def test_get_sessions_for_peer_with_empty_filter(client, sample_data):
|
||||
"""Test getting sessions for peer with empty filter object"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/sessions",
|
||||
json={"filter": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_create_and_get_messages_for_peer(client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create messages for the peer
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{
|
||||
"content": "Hello world",
|
||||
"peer_id": test_peer.name,
|
||||
"metadata": {"message_key": "message_value"},
|
||||
},
|
||||
{
|
||||
"content": "Second message",
|
||||
"peer_id": test_peer.name,
|
||||
"metadata": {"message_key": "message_value2"},
|
||||
},
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert len(data) == 2
|
||||
assert data[0]["content"] == "Hello world"
|
||||
assert data[1]["content"] == "Second message"
|
||||
assert data[0]["metadata"] == {"message_key": "message_value"}
|
||||
|
||||
# Get messages for the peer
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) == 2
|
||||
assert data["items"][0]["content"] == "Hello world"
|
||||
assert data["items"][1]["content"] == "Second message"
|
||||
assert data["items"][0]["metadata"] == {"message_key": "message_value"}
|
||||
assert data["items"][1]["metadata"] == {"message_key": "message_value2"}
|
||||
|
||||
|
||||
def test_get_messages_for_peer_with_reverse(client, sample_data):
|
||||
"""Test getting messages for peer with reverse parameter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Create messages
|
||||
client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{"content": "First message", "peer_id": test_peer.name},
|
||||
{"content": "Second message", "peer_id": test_peer.name},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
# Test normal order
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
normal_data = response.json()
|
||||
|
||||
# Test reversed order
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages/list?reverse=true",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
reversed_data = response.json()
|
||||
|
||||
# Both should have items
|
||||
assert len(normal_data["items"]) > 0
|
||||
assert len(reversed_data["items"]) > 0
|
||||
|
||||
|
||||
def test_get_messages_for_peer_with_empty_filter(client, sample_data):
|
||||
"""Test getting messages for peer with empty filter object"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages/list",
|
||||
json={"filter": {}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_get_messages_for_peer_with_null_filter(client, sample_data):
|
||||
"""Test getting messages for peer with null filter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages/list",
|
||||
json={"filter": None},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_chat(client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
target_peer = str(generate_nanoid())
|
||||
|
||||
# Test chat endpoint
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
json={
|
||||
"queries": "Hello, how are you?",
|
||||
"stream": False,
|
||||
"target": target_peer,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "content" in data
|
||||
|
||||
|
||||
def test_chat_with_optional_params(client, sample_data):
|
||||
"""Test chat endpoint with optional parameters"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
session_id = str(generate_nanoid())
|
||||
|
||||
# Create a session first
|
||||
client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_id, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
|
||||
# Test chat without optional parameters
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
json={
|
||||
"queries": "Hello, how are you?",
|
||||
"stream": False,
|
||||
"session_id": session_id,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "content" in data
|
||||
|
||||
|
||||
def test_get_peer_representation_with_session(client, sample_data):
|
||||
"""Test peer representation with session_id parameter"""
|
||||
test_workspace, test_peer = sample_data
|
||||
session_id = str(generate_nanoid())
|
||||
|
||||
# Create a session first
|
||||
client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_id, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
|
||||
# Test representation scoped to session
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/representation",
|
||||
json={
|
||||
"session_id": session_id,
|
||||
"queries": "Hello, how are you?",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_search_peer(client, sample_data):
|
||||
"""Test the peer search functionality"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Add some messages to search through
|
||||
client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{"content": "Search this content", "peer_id": test_peer.name},
|
||||
{"content": "Another searchable message", "peer_id": test_peer.name},
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
# Search with a query
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/search",
|
||||
json="search query",
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Response should have pagination structure
|
||||
assert "items" in data
|
||||
assert "total" in data
|
||||
assert "page" in data
|
||||
assert "size" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_search_peer_empty_query(client, sample_data):
|
||||
"""Test peer search with empty query"""
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
# Search with empty query
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/search", json=""
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Response should still have proper pagination structure
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_search_peer_nonexistent(client, sample_data):
|
||||
"""Test searching a peer that doesn't exist"""
|
||||
test_workspace, _ = sample_data
|
||||
nonexistent_peer_id = str(generate_nanoid())
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{nonexistent_peer_id}/search",
|
||||
json="test query",
|
||||
)
|
||||
# This should probably return 404 or handle gracefully
|
||||
# The exact behavior depends on the crud.search implementation
|
||||
assert response.status_code in [200, 404, 422]
|
||||
|
|
@ -3,11 +3,11 @@ from nanoid import generate as generate_nanoid
|
|||
from src.security import JWTParams, create_jwt
|
||||
|
||||
|
||||
def test_create_app_with_auth(auth_client):
|
||||
def test_create_workspace_with_auth(auth_client):
|
||||
name = str(generate_nanoid())
|
||||
|
||||
response = auth_client.post(
|
||||
"/v1/apps", json={"name": name, "metadata": {"key": "value"}}
|
||||
"/v1/workspaces", json={"name": name, "metadata": {"key": "value"}}
|
||||
)
|
||||
|
||||
# Check expected behavior based on auth type
|
||||
|
|
@ -26,7 +26,7 @@ def test_auth_response_time(auth_client):
|
|||
start_time = time.time()
|
||||
|
||||
response = auth_client.post(
|
||||
"/v1/apps", json={"name": name, "metadata": {"key": "value"}}
|
||||
"/v1/workspaces", json={"name": name, "metadata": {"key": "value"}}
|
||||
)
|
||||
|
||||
end_time = time.time()
|
||||
|
|
@ -43,18 +43,12 @@ def test_auth_response_time(auth_client):
|
|||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_get_or_create_app_with_auth(auth_client):
|
||||
def test_get_or_create_workspace_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}")
|
||||
response = auth_client.post(
|
||||
"/v1/workspaces", json={"name": name, "metadata": {"key": "value"}}
|
||||
)
|
||||
|
||||
if auth_client.auth_type != "admin":
|
||||
assert response.status_code == 401
|
||||
|
|
@ -63,223 +57,160 @@ def test_get_or_create_app_with_auth(auth_client):
|
|||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_get_app_by_id_with_auth(auth_client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
def test_get_workspace_with_auth(auth_client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the app_id in the JWT
|
||||
# For non-admin, include the workspace name in the JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
response = auth_client.get(f"/v1/apps?app_id={test_app.public_id}")
|
||||
response = auth_client.post("/v1/workspaces", json={"name": test_workspace.name})
|
||||
|
||||
# Admin JWT or JWT with matching app_id should be allowed
|
||||
# Admin JWT or JWT with matching workspace 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
|
||||
def test_update_workspace_with_auth(auth_client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the app_id in the JWT
|
||||
# For non-admin, include the workspace name 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))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
new_name = str(generate_nanoid())
|
||||
response = auth_client.put(
|
||||
f"/v1/apps/{test_app.public_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}",
|
||||
json={"name": new_name, "metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
|
||||
# Only admin JWT or JWT with matching app_id should be allowed
|
||||
# Only admin JWT or JWT with matching workspace 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
|
||||
def test_update_workspace_with_wrong_auth(auth_client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
different_app = str(generate_nanoid())
|
||||
different_workspace = str(generate_nanoid())
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the *wrong* app_id in the JWT
|
||||
# For non-admin, include the *wrong* workspace name in the JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(ap=different_app))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=different_workspace))}"
|
||||
)
|
||||
|
||||
new_name = str(generate_nanoid())
|
||||
response = auth_client.put(
|
||||
f"/v1/apps/{test_app.public_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}",
|
||||
json={"name": new_name, "metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
|
||||
# Only admin JWT or JWT with matching app_id should be allowed
|
||||
# Only admin JWT or JWT with matching workspace should be allowed
|
||||
if auth_client.auth_type == "admin":
|
||||
assert response.status_code == 200
|
||||
else:
|
||||
# wrong app_id should be rejected
|
||||
# wrong workspace should be rejected
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
def test_create_user_with_auth(auth_client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
def test_create_peer_with_auth(auth_client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the app_id in the JWT
|
||||
# For non-admin, include the workspace name in the JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
name = str(generate_nanoid())
|
||||
response = auth_client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users",
|
||||
json={"name": name, "metadata": {"user_key": "user_value"}},
|
||||
f"/v1/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": name, "metadata": {"peer_key": "peer_value"}},
|
||||
)
|
||||
|
||||
# Only admin JWT or JWT with matching app_id should be allowed
|
||||
# Only admin JWT or JWT with matching workspace 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
|
||||
def test_get_peer_by_name_with_auth(auth_client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the app_id in the JWT
|
||||
# For non-admin, include the workspace name in the JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
response = auth_client.get(
|
||||
f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}"
|
||||
# Use POST /list endpoint to get peers
|
||||
response = auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/list",
|
||||
json={"filter": {"name": test_peer.name}},
|
||||
)
|
||||
|
||||
# Admin JWT or JWT with matching app_id should be allowed
|
||||
# Admin JWT or JWT with matching workspace 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
|
||||
# Test with peer-scoped JWT
|
||||
if auth_client.auth_type == "empty":
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(p=test_peer.name))}"
|
||||
)
|
||||
|
||||
response = auth_client.get(
|
||||
f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}"
|
||||
# Get specific peer using get_or_create endpoint
|
||||
response = auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/peers", json={"name": test_peer.name}
|
||||
)
|
||||
|
||||
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
|
||||
def test_update_peer_with_auth(auth_client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the app_id in the JWT
|
||||
# For non-admin, include the workspace name 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))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
new_name = str(generate_nanoid())
|
||||
response = auth_client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={"name": new_name, "metadata": {"updated_key": "updated_value"}},
|
||||
)
|
||||
|
||||
# Admin JWT or JWT with matching app_id should be allowed
|
||||
# Admin JWT or JWT with matching workspace 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
|
||||
# Test with peer-scoped JWT
|
||||
if auth_client.auth_type == "empty":
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(p=test_peer.name, w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
response = auth_client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}",
|
||||
json={
|
||||
"name": str(generate_nanoid()),
|
||||
"metadata": {"user_key": "user_value"},
|
||||
"metadata": {"peer_key": "peer_value"},
|
||||
},
|
||||
)
|
||||
|
||||
|
|
@ -287,51 +218,54 @@ def test_update_user_with_auth(auth_client, sample_data):
|
|||
|
||||
|
||||
def test_create_session_with_auth(auth_client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
|
||||
if auth_client.auth_type == "empty":
|
||||
# For non-admin, include the app_id and user_id in the JWT
|
||||
# For non-admin, include the workspace name in the JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(ap=test_app.public_id, us=test_user.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
session_name = str(generate_nanoid())
|
||||
response = auth_client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": session_name, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
|
||||
# Only admin JWT or JWT with matching app_id and user_id should be allowed
|
||||
# Only admin JWT or JWT with matching workspace 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
|
||||
# Test with peer-scoped JWT
|
||||
if auth_client.auth_type == "empty":
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(p=test_peer.name))}"
|
||||
)
|
||||
|
||||
session_name2 = str(generate_nanoid())
|
||||
response = auth_client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": session_name2, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_get_session_by_id_with_auth(auth_client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
def test_get_session_by_name_with_auth(auth_client, sample_data):
|
||||
test_workspace, test_peer = 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))}"
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
session_name = str(generate_nanoid())
|
||||
create_response = auth_client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": session_name, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
|
||||
if auth_client.auth_type not in ["admin", "empty"]:
|
||||
|
|
@ -339,290 +273,65 @@ def test_get_session_by_id_with_auth(auth_client, sample_data):
|
|||
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}"
|
||||
# Test with workspace scoped JWT - get the same session
|
||||
response = auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions", json={"name": session_name}
|
||||
)
|
||||
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))}"
|
||||
f"Bearer {create_jwt(JWTParams(s=session_name))}"
|
||||
)
|
||||
|
||||
response = auth_client.get(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions?session_id={session_id}"
|
||||
response = auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": session_name},
|
||||
)
|
||||
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)
|
||||
# Test with wrong session_name (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()}"
|
||||
auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": generate_nanoid()},
|
||||
).status_code
|
||||
== 401
|
||||
)
|
||||
|
||||
# Test with user-scoped JWT
|
||||
# Test with peer-scoped JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(us=test_user.public_id))}"
|
||||
f"Bearer {create_jwt(JWTParams(p=test_peer.name))}"
|
||||
)
|
||||
|
||||
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}"
|
||||
auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": session_name},
|
||||
).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 workspace-scoped JWT
|
||||
auth_client.headers["Authorization"] = (
|
||||
f"Bearer {create_jwt(JWTParams(w=test_workspace.name))}"
|
||||
)
|
||||
|
||||
# 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()}"
|
||||
auth_client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"name": session_name},
|
||||
).status_code
|
||||
== 200
|
||||
)
|
||||
|
||||
# Test with wrong session_name using DELETE endpoint (should be 404 since session doesn't exist)
|
||||
wrong_session_name = generate_nanoid()
|
||||
assert (
|
||||
auth_client.delete(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{wrong_session_name}"
|
||||
).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
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -1,88 +0,0 @@
|
|||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_create_user(client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
name = str(generate_nanoid())
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users",
|
||||
json={"name": name, "metadata": {"user_key": "user_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == name
|
||||
assert data["metadata"] == {"user_key": "user_value"}
|
||||
assert "id" in 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?user_id={test_user.public_id}"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == test_user.name
|
||||
assert data["id"] == str(test_user.public_id)
|
||||
|
||||
|
||||
def test_get_user_by_name(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
response = client.get(f"/v1/apps/{test_app.public_id}/users/name/{test_user.name}")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == test_user.name
|
||||
assert data["id"] == str(test_user.public_id)
|
||||
|
||||
|
||||
def test_get_or_create_user(client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
name = str(generate_nanoid())
|
||||
response = client.get(f"/v1/apps/{test_app.public_id}/users/name/{name}")
|
||||
assert response.status_code == 404
|
||||
response = client.get(f"/v1/apps/{test_app.public_id}/users/get_or_create/{name}")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["name"] == name
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_users(client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"user_key": "user_value"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"user_key": "user_value"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users",
|
||||
json={"name": str(generate_nanoid()), "metadata": {"user_key": "user_value2"}},
|
||||
)
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/list",
|
||||
json={"filter": {"user_key": "user_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) > 0
|
||||
assert len(data["items"]) == 2
|
||||
assert data["items"][0]["metadata"] == {"user_key": "user_value"}
|
||||
|
||||
|
||||
def test_update_user(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
new_name = str(generate_nanoid())
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}",
|
||||
json={"name": new_name, "metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
print(new_name)
|
||||
print(data)
|
||||
assert data["name"] == new_name
|
||||
assert data["metadata"] == {"new_key": "new_value"}
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_app_validations_api(client):
|
||||
def test_workspace_validations_api(client):
|
||||
# Test name too short
|
||||
response = client.post("/v1/apps", json={"name": "", "metadata": {}})
|
||||
response = client.post("/v1/workspaces", json={"name": "", "metadata": {}})
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
|
|
@ -11,7 +11,7 @@ def test_app_validations_api(client):
|
|||
assert error["type"] == "string_too_short"
|
||||
|
||||
# Test name too long
|
||||
response = client.post("/v1/apps", json={"name": "a" * 101, "metadata": {}})
|
||||
response = client.post("/v1/workspaces", json={"name": "a" * 101, "metadata": {}})
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
|
|
@ -19,19 +19,21 @@ def test_app_validations_api(client):
|
|||
assert error["type"] == "string_too_long"
|
||||
|
||||
# Test invalid metadata type
|
||||
response = client.post("/v1/apps", json={"name": "test", "metadata": "not a dict"})
|
||||
response = client.post(
|
||||
"/v1/workspaces", json={"name": "test", "metadata": "not a dict"}
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "metadata"]
|
||||
assert error["type"] == "dict_type"
|
||||
|
||||
|
||||
def test_user_validations_api(client, sample_data):
|
||||
test_app, _ = sample_data
|
||||
def test_peer_validations_api(client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# Test name too short
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users", json={"name": "", "metadata": {}}
|
||||
f"/v1/workspaces/{test_workspace.name}/peers", json={"name": "", "metadata": {}}
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
|
|
@ -41,7 +43,8 @@ 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/workspaces/{test_workspace.name}/peers",
|
||||
json={"name": "a" * 101, "metadata": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
|
|
@ -51,312 +54,49 @@ def test_user_validations_api(client, sample_data):
|
|||
|
||||
|
||||
def test_message_validations_api(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
# Create a test session first
|
||||
session_id = str(generate_nanoid())
|
||||
session_response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={"metadata": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_id, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
session_id = session_response.json()["id"]
|
||||
assert session_response.status_code == 200
|
||||
|
||||
# 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": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{session_id}/messages",
|
||||
json={
|
||||
"messages": [
|
||||
{"content": "a" * 50001, "peer_id": test_peer.name, "metadata": {}}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "content"]
|
||||
assert "content" in str(error["loc"])
|
||||
assert error["msg"] == "String should have at most 50000 characters"
|
||||
assert error["type"] == "string_too_long"
|
||||
|
||||
# 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": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "is_user"]
|
||||
assert error["type"] == "bool_parsing"
|
||||
|
||||
|
||||
def test_collection_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
assert error["msg"] == "String should have at least 1 character"
|
||||
assert error["type"] == "string_too_short"
|
||||
|
||||
# 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": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
assert error["msg"] == "String should have at most 100 characters"
|
||||
assert error["type"] == "string_too_long"
|
||||
|
||||
# Test 'honcho' name restriction
|
||||
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
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
assert error["msg"] == "Value error, Collection name cannot be 'honcho'"
|
||||
assert error["type"] == "value_error"
|
||||
|
||||
|
||||
def test_document_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
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": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "content"]
|
||||
assert error["msg"] == "String should have at least 1 character"
|
||||
assert error["type"] == "string_too_short"
|
||||
|
||||
# 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": {}},
|
||||
)
|
||||
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"
|
||||
|
||||
|
||||
def test_document_query_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
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},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "query"]
|
||||
assert error["msg"] == "String should have at least 1 character"
|
||||
assert error["type"] == "string_too_short"
|
||||
|
||||
# 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},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "query"]
|
||||
assert error["msg"] == "String should have at most 1000 characters"
|
||||
assert error["type"] == "string_too_long"
|
||||
|
||||
# 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},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "top_k"]
|
||||
assert error["msg"] == "Input should be greater than or equal to 1"
|
||||
assert error["type"] == "greater_than_equal"
|
||||
|
||||
# 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},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "top_k"]
|
||||
assert error["msg"] == "Input should be less than or equal to 50"
|
||||
assert error["type"] == "less_than_equal"
|
||||
|
||||
|
||||
def test_message_batch_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
session_id = session_response.json()["id"]
|
||||
|
||||
# Test batch too large
|
||||
messages = [
|
||||
{"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},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "messages"]
|
||||
assert "List should have at most 100 items after validation" in error["msg"]
|
||||
assert error["type"] == "too_long"
|
||||
|
||||
|
||||
def test_metamessage_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
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": {}},
|
||||
)
|
||||
message_id = message_response.json()["id"]
|
||||
|
||||
# Test label too short
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
|
||||
json={
|
||||
"label": "",
|
||||
"content": "test content",
|
||||
"session_id": session_id,
|
||||
"message_id": message_id,
|
||||
"metadata": {},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "label"]
|
||||
assert error["msg"] == "String should have at least 1 character"
|
||||
assert error["type"] == "string_too_short"
|
||||
|
||||
# Test label too long
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
|
||||
json={
|
||||
"label": "a" * 51,
|
||||
"content": "test content",
|
||||
"session_id": session_id,
|
||||
"message_id": message_id,
|
||||
"metadata": {},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "label"]
|
||||
assert error["msg"] == "String should have at most 50 characters"
|
||||
assert error["type"] == "string_too_long"
|
||||
|
||||
# Test content too long
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/metamessages",
|
||||
json={
|
||||
"label": "test_type",
|
||||
"content": "a" * 50001,
|
||||
"message_id": message_id,
|
||||
"session_id": session_id,
|
||||
"metadata": {},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "content"]
|
||||
assert error["msg"] == "String should have at most 50000 characters"
|
||||
assert error["type"] == "string_too_long"
|
||||
|
||||
|
||||
def test_collection_update_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
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": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
assert error["msg"] == "Value error, Collection name cannot be 'honcho'"
|
||||
assert error["type"] == "value_error"
|
||||
|
||||
|
||||
def test_document_update_validations_api(client, sample_data):
|
||||
test_app, test_user = 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": {}},
|
||||
)
|
||||
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": {}},
|
||||
)
|
||||
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": {}},
|
||||
)
|
||||
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"
|
||||
|
||||
|
||||
def test_session_validations_api(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
# Create a test session first
|
||||
session_id = str(generate_nanoid())
|
||||
session_response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={"metadata": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={
|
||||
"id": session_id,
|
||||
"peer_names": {test_peer.name: {}},
|
||||
"metadata": {"test_key": "test_value"},
|
||||
"configuration": {"test_flag": "test_value"},
|
||||
},
|
||||
)
|
||||
session_id = session_response.json()["id"]
|
||||
assert session_response.status_code == 200
|
||||
|
||||
# Test invalid metadata type
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{session_id}",
|
||||
json={"metadata": "not a dict"},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
|
@ -365,91 +105,49 @@ def test_session_validations_api(client, sample_data):
|
|||
assert error["type"] == "dict_type"
|
||||
|
||||
# Test empty update
|
||||
# This should work but not change the session's metadata or configuration
|
||||
response = client.put(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions/{session_id}",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{session_id}",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
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,
|
||||
# Test that the session's metadata and configuration are not changed
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={
|
||||
"id": session_id,
|
||||
},
|
||||
)
|
||||
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)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == {"test_key": "test_value"}
|
||||
assert data["configuration"] == {"test_flag": "test_value"}
|
||||
|
||||
test_app, test_user = sample_data
|
||||
# Create a session first since agent queries are likely session-based
|
||||
|
||||
def test_agent_query_validations_api(client, sample_data):
|
||||
test_workspace, test_peer = sample_data
|
||||
# Create a session first since agent queries are session-based
|
||||
session_id = str(generate_nanoid())
|
||||
session_response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={"metadata": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_id, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
session_id = session_response.json()["id"]
|
||||
assert session_response.status_code == 200
|
||||
|
||||
# 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},
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
params={"session_id": session_id, "target": "test_target"},
|
||||
json={"queries": "a" * 9999, "stream": False},
|
||||
)
|
||||
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},
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
params={"session_id": session_id, "target": "test_target"},
|
||||
json={"queries": "a" * 10001, "stream": False},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
|
|
@ -459,15 +157,17 @@ def test_agent_query_validations_api(client, sample_data, monkeypatch):
|
|||
|
||||
# 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)]},
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
params={"session_id": session_id, "target": "test_target"},
|
||||
json={"queries": ["a" * 9999 for _ in range(25)], "stream": False},
|
||||
)
|
||||
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)]},
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
params={"session_id": session_id, "target": "test_target"},
|
||||
json={"queries": ["test" for _ in range(26)], "stream": False},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
|
|
@ -476,8 +176,9 @@ def test_agent_query_validations_api(client, sample_data, monkeypatch):
|
|||
|
||||
# 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]},
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
params={"session_id": session_id, "target": "test_target"},
|
||||
json={"queries": ["a" * 10001], "stream": False},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
|
|
@ -487,73 +188,56 @@ def test_agent_query_validations_api(client, sample_data, monkeypatch):
|
|||
|
||||
# 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
|
||||
f"/v1/workspaces/{test_workspace.name}/peers/{test_peer.name}/chat",
|
||||
params={"session_id": session_id, "target": "test_target"},
|
||||
json={"queries": "a" * 100, "stream": False}, # 100 chars should be fine
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_required_field_validations_api(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
session_id = str(generate_nanoid())
|
||||
session_response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={"metadata": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_id, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
session_id = session_response.json()["id"]
|
||||
assert session_response.status_code == 200
|
||||
|
||||
# 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": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{session_id}/messages",
|
||||
json={"messages": [{"peer_id": test_peer.name, "metadata": {}}]},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "content"]
|
||||
assert "content" in str(error["loc"])
|
||||
assert error["type"] == "missing"
|
||||
|
||||
# Test missing required is_user in message
|
||||
# Test missing required peer_id 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": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{session_id}/messages",
|
||||
json={"messages": [{"content": "test", "metadata": {}}]},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "is_user"]
|
||||
assert error["type"] == "missing"
|
||||
|
||||
# Test missing required name in collection
|
||||
response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/collections",
|
||||
json={"metadata": {}},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "name"]
|
||||
assert "peer_id" in str(error["loc"]) or "peer_name" in str(error["loc"])
|
||||
assert error["type"] == "missing"
|
||||
|
||||
|
||||
def test_filter_validations_api(client, sample_data):
|
||||
test_app, test_user = sample_data
|
||||
test_workspace, test_peer = sample_data
|
||||
# Create a session first
|
||||
session_id = str(generate_nanoid())
|
||||
session_response = client.post(
|
||||
f"/v1/apps/{test_app.public_id}/users/{test_user.public_id}/sessions",
|
||||
json={"metadata": {}},
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions",
|
||||
json={"id": session_id, "peer_names": {test_peer.name: {}}},
|
||||
)
|
||||
session_id = session_response.json()["id"]
|
||||
assert session_response.status_code == 200
|
||||
|
||||
# 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"},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
error = response.json()["detail"][0]
|
||||
assert error["loc"] == ["body", "filter"]
|
||||
assert error["type"] == "dict_type"
|
||||
|
||||
# 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",
|
||||
f"/v1/workspaces/{test_workspace.name}/sessions/{session_id}/messages/list",
|
||||
json={"filter": "not a dict"},
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
|
|
|||
|
|
@ -0,0 +1,248 @@
|
|||
import pytest
|
||||
from nanoid import generate as generate_nanoid
|
||||
|
||||
|
||||
def test_get_or_create_workspace(client):
|
||||
name = str(generate_nanoid())
|
||||
|
||||
# This should create the workspace using POST /v1/workspaces
|
||||
response = client.post("/v1/workspaces", json={"name": name})
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
assert "id" in data
|
||||
|
||||
|
||||
def test_get_or_create_workspace_with_configuration(client):
|
||||
"""Test workspace creation with configuration parameter"""
|
||||
name = str(generate_nanoid())
|
||||
configuration = {"feature1": True, "feature2": False}
|
||||
|
||||
response = client.post(
|
||||
"/v1/workspaces", json={"name": name, "configuration": configuration}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_get_or_create_workspace_with_all_optional_params(client):
|
||||
"""Test workspace creation with all optional parameters"""
|
||||
name = str(generate_nanoid())
|
||||
metadata = {"key": "value", "number": 42}
|
||||
configuration = {"experimental": True, "beta": False}
|
||||
|
||||
response = client.post(
|
||||
"/v1/workspaces",
|
||||
json={"name": name, "metadata": metadata, "configuration": configuration},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
assert data["metadata"] == metadata
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_get_or_create_existing_workspace(client):
|
||||
name = str(generate_nanoid())
|
||||
|
||||
# Create the workspace
|
||||
response = client.post(
|
||||
"/v1/workspaces", json={"name": name, "metadata": {"key": "value"}}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
workspace1 = response.json()
|
||||
|
||||
# Try to create the same workspace again - should return existing workspace
|
||||
response = client.post(
|
||||
"/v1/workspaces", json={"name": name, "metadata": {"key": "value"}}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
workspace2 = response.json()
|
||||
|
||||
# Both should be the same workspace
|
||||
assert workspace1["id"] == workspace2["id"]
|
||||
assert workspace1["metadata"] == workspace2["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_workspaces(client, db_session, sample_data):
|
||||
# create a test workspace with metadata
|
||||
response = client.post(
|
||||
"/v1/workspaces",
|
||||
json={
|
||||
"name": "test_workspace",
|
||||
"metadata": {"test_key": "test_value"},
|
||||
},
|
||||
)
|
||||
|
||||
response = client.post(
|
||||
"/v1/workspaces/list",
|
||||
json={},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert len(data["items"]) > 0
|
||||
|
||||
response = client.post(
|
||||
"/v1/workspaces/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"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_workspaces_with_empty_filter(client, db_session, sample_data):
|
||||
"""Test workspace listing with empty filter object"""
|
||||
response = client.post("/v1/workspaces/list", json={"filter": {}})
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_workspaces_with_null_filter(client, db_session, sample_data):
|
||||
"""Test workspace listing with null filter"""
|
||||
response = client.post("/v1/workspaces/list", json={"filter": None})
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_update_workspace(client, sample_data):
|
||||
test_workspace, _ = sample_data
|
||||
_new_name = str(generate_nanoid())
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}",
|
||||
json={"metadata": {"new_key": "new_value"}},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == {"new_key": "new_value"}
|
||||
|
||||
|
||||
def test_update_workspace_with_configuration(client, sample_data):
|
||||
"""Test workspace update with configuration parameter"""
|
||||
test_workspace, _ = sample_data
|
||||
configuration = {"new_feature": True, "legacy_feature": False}
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}", json={"configuration": configuration}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_update_workspace_with_all_optional_params(client, sample_data):
|
||||
"""Test workspace update with both metadata and configuration"""
|
||||
test_workspace, _ = sample_data
|
||||
metadata = {"updated_key": "updated_value", "count": 100}
|
||||
configuration = {"experimental": True, "beta": True}
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}",
|
||||
json={"metadata": metadata, "configuration": configuration},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["metadata"] == metadata
|
||||
assert data["configuration"] == configuration
|
||||
|
||||
|
||||
def test_update_workspace_with_null_metadata(client, sample_data):
|
||||
"""Test workspace update with null metadata (should clear metadata)"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# First set some metadata
|
||||
client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}", json={"metadata": {"temp": "value"}}
|
||||
)
|
||||
|
||||
# Then clear it with null
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}", json={"metadata": None}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
# The behavior might be to keep existing metadata or clear it -
|
||||
# adjust this assertion based on actual behavior
|
||||
assert "metadata" in data
|
||||
|
||||
|
||||
def test_update_workspace_with_null_configuration(client, sample_data):
|
||||
"""Test workspace update with null configuration"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
response = client.put(
|
||||
f"/v1/workspaces/{test_workspace.name}", json={"configuration": None}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert "configuration" in data
|
||||
|
||||
|
||||
def test_create_duplicate_workspace_name(client):
|
||||
# Create an workspace
|
||||
name = str(generate_nanoid())
|
||||
response = client.post("/v1/workspaces", json={"name": name})
|
||||
assert response.status_code == 200
|
||||
|
||||
# Try to create another workspace with the same name - should return existing workspace
|
||||
response = client.post("/v1/workspaces", json={"name": name})
|
||||
|
||||
# Should return the existing workspace with 200 status (get_or_create behavior)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == name
|
||||
|
||||
|
||||
def test_search_workspace(client, sample_data):
|
||||
"""Test the workspace search functionality"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# Test search with a query
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{test_workspace.name}/search", json="test search query"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Response should have pagination structure
|
||||
assert "items" in data
|
||||
assert "total" in data
|
||||
assert "page" in data
|
||||
assert "size" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_search_workspace_empty_query(client, sample_data):
|
||||
"""Test the workspace search with empty query"""
|
||||
test_workspace, _ = sample_data
|
||||
|
||||
# Test search with empty query
|
||||
response = client.post(f"/v1/workspaces/{test_workspace.name}/search", json="")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
# Response should still have proper pagination structure
|
||||
assert "items" in data
|
||||
assert isinstance(data["items"], list)
|
||||
|
||||
|
||||
def test_search_workspace_nonexistent(client):
|
||||
"""Test searching a workspace that doesn't exist"""
|
||||
nonexistent_workspace_id = str(generate_nanoid())
|
||||
|
||||
response = client.post(
|
||||
f"/v1/workspaces/{nonexistent_workspace_id}/search", json="test query"
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
|
@ -2,153 +2,90 @@ import pytest
|
|||
from pydantic import ValidationError
|
||||
|
||||
from src.schemas import (
|
||||
AppCreate,
|
||||
CollectionCreate,
|
||||
DocumentCreate,
|
||||
DocumentQuery,
|
||||
MessageBatchCreate,
|
||||
MessageCreate,
|
||||
MetamessageCreate,
|
||||
UserCreate,
|
||||
PeerCreate,
|
||||
SessionCreate,
|
||||
WorkspaceCreate,
|
||||
)
|
||||
|
||||
|
||||
class TestAppValidations:
|
||||
def test_valid_app_create(self):
|
||||
app = AppCreate(name="test", metadata={})
|
||||
assert app.name == "test"
|
||||
assert app.metadata == {}
|
||||
class TestWorkspaceValidations:
|
||||
def test_valid_workspace_create(self):
|
||||
workspace = WorkspaceCreate(name="test", metadata={})
|
||||
assert workspace.name == "test"
|
||||
assert workspace.metadata == {}
|
||||
|
||||
def test_app_name_too_short(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
AppCreate(name="", metadata={})
|
||||
WorkspaceCreate(name="", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_short"
|
||||
|
||||
def test_app_name_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
AppCreate(name="a" * 101, metadata={})
|
||||
WorkspaceCreate(name="a" * 101, metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
def test_app_invalid_metadata_type(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
AppCreate(name="test", metadata="not a dict")
|
||||
WorkspaceCreate(name="test", metadata="not a dict")
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "dict_type"
|
||||
|
||||
|
||||
class TestUserValidations:
|
||||
def test_valid_user_create(self):
|
||||
user = UserCreate(name="test", metadata={})
|
||||
assert user.name == "test"
|
||||
assert user.metadata == {}
|
||||
class TestPeerValidations:
|
||||
def test_valid_peer_create(self):
|
||||
peer = PeerCreate(name="test", metadata={})
|
||||
assert peer.name == "test"
|
||||
assert peer.metadata == {}
|
||||
|
||||
def test_user_name_too_short(self):
|
||||
def test_valid_peer_create_configuration(self):
|
||||
peer = PeerCreate(name="test", metadata={}, configuration={"test": True})
|
||||
assert peer.name == "test"
|
||||
assert peer.metadata == {}
|
||||
assert peer.configuration == {"test": True}
|
||||
|
||||
def test_peer_name_too_short(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
UserCreate(name="", metadata={})
|
||||
PeerCreate(name="", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_short"
|
||||
|
||||
def test_user_name_too_long(self):
|
||||
def test_peer_name_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
UserCreate(name="a" * 101, metadata={})
|
||||
PeerCreate(name="a" * 101, metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
|
||||
class TestSessionValidations:
|
||||
def test_valid_session_create(self):
|
||||
session = SessionCreate(name="test", metadata={})
|
||||
assert session.name == "test"
|
||||
assert session.metadata == {}
|
||||
|
||||
def test_session_name_too_short(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
SessionCreate(name="", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_short"
|
||||
|
||||
|
||||
class TestMessageValidations:
|
||||
def test_valid_message_create(self):
|
||||
msg = MessageCreate(content="test", is_user=True, metadata={})
|
||||
msg = MessageCreate(content="test", peer_id="12345", metadata={})
|
||||
assert msg.content == "test"
|
||||
assert msg.is_user is True
|
||||
assert msg.peer_name == "12345"
|
||||
assert msg.metadata == {}
|
||||
|
||||
def test_message_content_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
MessageCreate(content="a" * 50001, is_user=True, metadata={})
|
||||
MessageCreate(content="a" * 50001, peer_id="12345", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
def test_message_invalid_is_user_type(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
MessageCreate(content="test", is_user="not a bool", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "bool_parsing"
|
||||
|
||||
|
||||
class TestMetamessageValidations:
|
||||
def test_valid_metamessage_create(self):
|
||||
meta = MetamessageCreate(
|
||||
label="test",
|
||||
content="test content",
|
||||
message_id="123",
|
||||
metadata={},
|
||||
)
|
||||
assert meta.label == "test"
|
||||
assert meta.content == "test content"
|
||||
assert meta.message_id == "123"
|
||||
|
||||
def test_label_too_short(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
MetamessageCreate(
|
||||
label="",
|
||||
content="test",
|
||||
message_id="123",
|
||||
metadata={},
|
||||
)
|
||||
error_dict = exc_info.value.errors(include_input=False)[0]
|
||||
assert error_dict["type"] == "string_too_short"
|
||||
assert error_dict["loc"] == ('label',)
|
||||
|
||||
def test_label_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
MetamessageCreate(
|
||||
label="a" * 51,
|
||||
content="test",
|
||||
message_id="123",
|
||||
metadata={},
|
||||
)
|
||||
error_dict = exc_info.value.errors(include_input=False)[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
assert error_dict["loc"] == ('label',)
|
||||
|
||||
def test_metamessage_content_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
MetamessageCreate(
|
||||
label="test",
|
||||
content="a" * 50001,
|
||||
message_id="123",
|
||||
metadata={},
|
||||
)
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
|
||||
class TestCollectionValidations:
|
||||
def test_valid_collection_create(self):
|
||||
collection = CollectionCreate(name="test", metadata={})
|
||||
assert collection.name == "test"
|
||||
assert collection.metadata == {}
|
||||
|
||||
def test_collection_name_too_short(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
CollectionCreate(name="", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_short"
|
||||
|
||||
def test_collection_name_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
CollectionCreate(name="a" * 101, metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
def test_collection_name_honcho(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
CollectionCreate(name="honcho", metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "value_error"
|
||||
|
||||
|
||||
class TestDocumentValidations:
|
||||
def test_valid_document_create(self):
|
||||
|
|
@ -167,53 +104,3 @@ class TestDocumentValidations:
|
|||
DocumentCreate(content="a" * 100001, metadata={})
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
|
||||
class TestDocumentQueryValidations:
|
||||
def test_valid_document_query(self):
|
||||
query = DocumentQuery(query="test query", top_k=5)
|
||||
assert query.query == "test query"
|
||||
assert query.top_k == 5
|
||||
|
||||
def test_query_too_short(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
DocumentQuery(query="", top_k=5)
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_short"
|
||||
|
||||
def test_query_too_long(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
DocumentQuery(query="a" * 1001, top_k=5)
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "string_too_long"
|
||||
|
||||
def test_top_k_too_small(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
DocumentQuery(query="test", top_k=0)
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "greater_than_equal"
|
||||
|
||||
def test_top_k_too_large(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
DocumentQuery(query="test", top_k=51)
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "less_than_equal"
|
||||
|
||||
|
||||
class TestMessageBatchValidations:
|
||||
def test_valid_message_batch(self):
|
||||
batch = MessageBatchCreate(
|
||||
messages=[MessageCreate(content="test", is_user=True, metadata={})]
|
||||
)
|
||||
assert len(batch.messages) == 1
|
||||
|
||||
def test_message_batch_too_large(self):
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
MessageBatchCreate(
|
||||
messages=[
|
||||
MessageCreate(content="test", is_user=True, metadata={})
|
||||
for _ in range(101)
|
||||
]
|
||||
)
|
||||
error_dict = exc_info.value.errors()[0]
|
||||
assert error_dict["type"] == "too_long"
|
||||
|
|
|
|||
|
|
@ -101,9 +101,11 @@ def test_unsupported_provider_initialization():
|
|||
|
||||
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()
|
||||
with (
|
||||
patch.dict(os.environ, {}, clear=True),
|
||||
pytest.raises(ValueError, match="API key is required"),
|
||||
):
|
||||
ModelClient()
|
||||
|
||||
|
||||
# Message Creation Tests
|
||||
|
|
|
|||
Loading…
Reference in New Issue