Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Eléonore Voisin
|
921a8fc845 | ||
|
Eléonore Voisin
|
0790b006ab | ||
|
Eléonore Voisin
|
f6da35e064 | ||
|
Eléonore Voisin
|
80248111d2 | ||
|
Eléonore Voisin
|
6b5826fc17 | ||
|
Eléonore Voisin
|
4c34608099 | ||
|
Eléonore Voisin
|
f7a60ff443 | ||
|
Eléonore Voisin
|
2294122d96 |
@@ -9,6 +9,3 @@ DB_PASSWORD=pass
|
||||
DB_PORT=15432
|
||||
|
||||
STORAGES_STATICFILES_BACKEND=django.contrib.staticfiles.storage.StaticFilesStorage
|
||||
AWS_S3_ENDPOINT_URL=http://localhost:9000
|
||||
AWS_S3_ACCESS_KEY_ID=conversations
|
||||
AWS_S3_SECRET_ACCESS_KEY=password
|
||||
|
||||
@@ -131,9 +131,6 @@ jobs:
|
||||
DB_PASSWORD: pass
|
||||
DB_PORT: 5432
|
||||
STORAGES_STATICFILES_BACKEND: django.contrib.staticfiles.storage.StaticFilesStorage
|
||||
AWS_S3_ENDPOINT_URL: http://localhost:9000
|
||||
AWS_S3_ACCESS_KEY_ID: conversations
|
||||
AWS_S3_SECRET_ACCESS_KEY: password
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
@@ -152,34 +149,11 @@ jobs:
|
||||
key: mail-templates-${{ hashFiles('src/mail/mjml') }}
|
||||
fail-on-cache-miss: true
|
||||
|
||||
- name: Start MinIO
|
||||
run: |
|
||||
docker pull minio/minio
|
||||
docker run -d --name minio \
|
||||
-p 9000:9000 \
|
||||
-e "MINIO_ACCESS_KEY=conversations" \
|
||||
-e "MINIO_SECRET_KEY=password" \
|
||||
-v /data/media:/data \
|
||||
minio/minio server --console-address :9001 /data
|
||||
|
||||
# Tool to wait for a service to be ready
|
||||
- name: Install Dockerize
|
||||
run: |
|
||||
curl -sSL https://github.com/jwilder/dockerize/releases/download/v0.8.0/dockerize-linux-amd64-v0.8.0.tar.gz | sudo tar -C /usr/local/bin -xzv
|
||||
|
||||
- name: Wait for MinIO to be ready
|
||||
run: |
|
||||
dockerize -wait tcp://localhost:9000 -timeout 10s
|
||||
|
||||
- name: Configure MinIO
|
||||
run: |
|
||||
MINIO=$(docker ps | grep minio/minio | sed -E 's/.*\s+([a-zA-Z0-9_-]+)$/\1/')
|
||||
docker exec ${MINIO} sh -c \
|
||||
"mc alias set conversations http://localhost:9000 conversations password && \
|
||||
mc alias ls && \
|
||||
mc mb conversations/conversations-media-storage && \
|
||||
mc version enable conversations/conversations-media-storage"
|
||||
|
||||
- name: Install Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
@@ -200,22 +174,3 @@ jobs:
|
||||
|
||||
- name: Run tests
|
||||
run: ~/.local/bin/pytest -n 2
|
||||
|
||||
security-trivy-critical:
|
||||
permissions:
|
||||
contents: read
|
||||
security-events: write
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Run Trivy analysis for critical vulnerabilities
|
||||
# We use main branch while we might still iterate on the action
|
||||
uses: numerique-gouv/action-trivy-cache/security-trivy-critical@main
|
||||
|
||||
security-trivy:
|
||||
permissions:
|
||||
contents: read
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Run Trivy analysis for vulnerabilities
|
||||
# We use main branch while we might still iterate on the action
|
||||
uses: numerique-gouv/action-trivy-cache/security-trivy@main
|
||||
|
||||
@@ -47,7 +47,6 @@ jobs:
|
||||
docker-image-name: 'docker.io/lasuite/conversations-backend:${{ github.sha }}'
|
||||
-
|
||||
name: Build and push
|
||||
if: always()
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
@@ -87,7 +86,6 @@ jobs:
|
||||
docker-image-name: 'docker.io/lasuite/conversations-frontend:${{ github.sha }}'
|
||||
-
|
||||
name: Build and push
|
||||
if: always()
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
|
||||
+4
-130
@@ -8,125 +8,10 @@ and this project adheres to
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Changed
|
||||
- 🐛(front) optimize chat
|
||||
- 📦️(front) update react
|
||||
- ✨(chat) Generate and edit conversation title
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🐛(e2e) fix test-e2e-chronium
|
||||
|
||||
## [0.0.10] - 2025-12-15
|
||||
|
||||
### Added
|
||||
|
||||
- ✨(front) add retry button
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🐛(front) fix long user messages
|
||||
- 🐛(front) fix "Maximum update depth exceeded" error in Chat component
|
||||
- 🐛(front) fix parsing documents display
|
||||
- 🐛(front) fix opacity input in error
|
||||
- 🐛(front) resolve React hydration errors
|
||||
- 🚑️(user) allow longer short names #182
|
||||
|
||||
## [0.0.9] - 2025-11-17
|
||||
|
||||
### Added
|
||||
- ✨(front) add code copy button
|
||||
- ✨(RAG) add generic collection RAG tools #159
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🔊(langfuse) enable tracing with redacted content #162
|
||||
|
||||
|
||||
## [0.0.8] - 2025-11-10
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🦺(front) Fix send prohibited file types
|
||||
- 🐛(front) fix target blank links in chat #103
|
||||
- 🚑️(posthog) pass str instead of UUID for user PK #134
|
||||
- ⚡️(web-search) keep running when tool call fails #137
|
||||
- ✨(summarize): new summarize tool integration #78
|
||||
|
||||
### Removed
|
||||
|
||||
- 🔥(posthog) remove posthog middleware for async mode fix #146
|
||||
|
||||
|
||||
## [0.0.7] - 2025-10-28
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🚑️(posthog) fix the posthog middleware for async mode #133
|
||||
|
||||
|
||||
## [0.0.6] - 2025-10-28
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🚑️(stats) fix tracking id in upload event #130
|
||||
|
||||
|
||||
## [0.0.5] - 2025-10-27
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🚑️(drag-drop) fix the rejection display on Safari #127
|
||||
|
||||
|
||||
## [0.0.4] - 2025-10-27
|
||||
|
||||
### Added
|
||||
|
||||
- ♿️(a11y) improve accessibility #135
|
||||
- 🌐(i18n) add dutch language #117
|
||||
|
||||
### Changed
|
||||
|
||||
- ⚡️(asgi) use `uvicorn` to serve backend #121
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🐛(front) fix mobile source
|
||||
- 🐛(attachments) reject the whole drag&drop if unsupported formats #123
|
||||
|
||||
|
||||
## [0.0.3] - 2025-10-21
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🚑️(web-search) fix missing argument in RAG backend #116
|
||||
|
||||
|
||||
## [0.0.2] - 2025-10-21
|
||||
|
||||
### Added
|
||||
|
||||
- ✨(front) add drag'n drop file
|
||||
- ✨(activation-codes) register users also on Brevo #98
|
||||
- 📈(posthog) add `sub` field to tracking #95
|
||||
|
||||
### Changed
|
||||
- 🔧(front) change links feedback tchap + settings popup
|
||||
- 🐛(front) code activation fix session end #93
|
||||
- 💬(wording) error page wording #102
|
||||
- ⚡️(web-search) allow to override returned chunks #107
|
||||
- 🐛(activation-codes) create contact in brevo before add to list #108
|
||||
- ⚗️(summarization) add system prompt to handle tool #112
|
||||
|
||||
|
||||
## [0.0.1] - 2025-10-19
|
||||
|
||||
### Changed
|
||||
|
||||
- 🎨(front) activation page footer
|
||||
- 👷(front) change size small modal
|
||||
- 🎨(front) retour ui global
|
||||
|
||||
### Changed
|
||||
|
||||
- 👷(front) fix button scrollDown
|
||||
- 💥(front) disable input when error occurred
|
||||
- 👷(front) fix scroll
|
||||
@@ -155,7 +40,6 @@ and this project adheres to
|
||||
- 🐛(chat) UI messages must have a unique identifier #6
|
||||
- ✨(llm) allow configuration from JSON file #22
|
||||
- 💥(agent) replace routing w/ tool calls #40
|
||||
- 🧱(storage) upload the user documents into S3 #86
|
||||
|
||||
### Added
|
||||
|
||||
@@ -177,14 +61,4 @@ and this project adheres to
|
||||
- 💄(chat) add code highlighting for LLM responses #67
|
||||
|
||||
|
||||
[unreleased]: https://github.com/suitenumerique/conversations/compare/v0.0.10...main
|
||||
[0.0.10]: https://github.com/suitenumerique/conversations/releases/v0.0.10
|
||||
[0.0.9]: https://github.com/suitenumerique/conversations/releases/v0.0.9
|
||||
[0.0.8]: https://github.com/suitenumerique/conversations/releases/v0.0.8
|
||||
[0.0.7]: https://github.com/suitenumerique/conversations/releases/v0.0.7
|
||||
[0.0.6]: https://github.com/suitenumerique/conversations/releases/v0.0.6
|
||||
[0.0.5]: https://github.com/suitenumerique/conversations/releases/v0.0.5
|
||||
[0.0.4]: https://github.com/suitenumerique/conversations/releases/v0.0.4
|
||||
[0.0.3]: https://github.com/suitenumerique/conversations/releases/v0.0.3
|
||||
[0.0.2]: https://github.com/suitenumerique/conversations/releases/v0.0.2
|
||||
[0.0.1]: https://github.com/suitenumerique/conversations/releases/v0.0.1
|
||||
[unreleased]: https://github.com/numerique-gouv/conversations/compare/HEAD...main
|
||||
|
||||
+3
-16
@@ -144,7 +144,7 @@ RUN rm -rf /var/cache/apk/*
|
||||
|
||||
ARG CONVERSATIONS_STATIC_ROOT=/data/static
|
||||
|
||||
# Gunicorn - not used by default but configuration file is provided
|
||||
# Gunicorn
|
||||
RUN mkdir -p /usr/local/etc/gunicorn
|
||||
COPY docker/files/usr/local/etc/gunicorn/conversations.py /usr/local/etc/gunicorn/conversations.py
|
||||
|
||||
@@ -158,18 +158,5 @@ COPY --from=link-collector ${CONVERSATIONS_STATIC_ROOT} ${CONVERSATIONS_STATIC_R
|
||||
# Copy conversations mails
|
||||
COPY --from=mail-builder /mail/backend/core/templates/mail /app/core/templates/mail
|
||||
|
||||
# The default command runs uvicorn ASGI server in conversations's main module
|
||||
# WEB_CONCURRENCY: number of workers to run <=> --workers=4
|
||||
ENV WEB_CONCURRENCY=4
|
||||
CMD [\
|
||||
"uvicorn",\
|
||||
"--app-dir=/app",\
|
||||
"--host=0.0.0.0",\
|
||||
"--timeout-graceful-shutdown=300",\
|
||||
"--limit-max-requests=20000",\
|
||||
"--lifespan=off",\
|
||||
"conversations.asgi:application"\
|
||||
]
|
||||
|
||||
# To run using gunicorn WSGI server use this instead:
|
||||
#CMD ["gunicorn", "-c", "/usr/local/etc/gunicorn/conversations.py", "conversations.wsgi:application"]
|
||||
# The default command runs gunicorn WSGI server in conversations's main module
|
||||
CMD ["gunicorn", "-c", "/usr/local/etc/gunicorn/conversations.py", "conversations.wsgi:application"]
|
||||
|
||||
@@ -126,7 +126,7 @@ build-frontend: ## build the frontend container
|
||||
build-e2e: cache ?=
|
||||
build-e2e: ## build the e2e container
|
||||
@$(MAKE) build-backend cache=$(cache)
|
||||
@$(COMPOSE_E2E) build frontend openmockllm-mistral $(cache)
|
||||
@$(COMPOSE_E2E) build frontend $(cache)
|
||||
.PHONY: build-e2e
|
||||
|
||||
down: ## stop and remove containers, networks, images, and volumes
|
||||
@@ -158,7 +158,7 @@ create-compose-with-models: ## override the docker-compose file with models
|
||||
run-e2e: ## start the e2e server
|
||||
run-e2e:
|
||||
@$(MAKE) run-backend
|
||||
@$(COMPOSE_E2E) up --force-recreate -d frontend openmockllm-mistral
|
||||
@$(COMPOSE_E2E) up --force-recreate -d frontend
|
||||
.PHONY: run-e2e
|
||||
|
||||
status: ## an alias for "docker compose ps"
|
||||
|
||||
@@ -115,31 +115,6 @@ To start all the services, except the frontend container, you can use the follow
|
||||
$ make run-backend
|
||||
```
|
||||
|
||||
**Setup a basic LLM call**
|
||||
|
||||
To be able to use Conversations, you need to configure at least one Large Language Model (LLM) provider.
|
||||
You can do so by setting the appropriate environment variables in the `env.d/development/common` file:
|
||||
|
||||
```ini
|
||||
AI_BASE_URL=http://host.docker.internal:12434/v1/
|
||||
AI_MODEL=gemma3:4b
|
||||
AI_API_KEY=XXX
|
||||
```
|
||||
|
||||
for a local ollama, or by running a local LLM with docker-compose:
|
||||
|
||||
```shellscript
|
||||
$ make create-compose-with-models
|
||||
```
|
||||
|
||||
which will create a `compose.override.yml` file to start a local models `ai/smollm2`
|
||||
which can be changed later by editing the `compose.override.yml` file.
|
||||
|
||||
You will need to call `make run` after changing the `env.d/development/common`
|
||||
or `compose.override.yml` file.
|
||||
|
||||
You can find more information about configuring LLM providers in the [LLM Configuration](docs/llm-configuration.md) documentation.
|
||||
|
||||
**Adding content**
|
||||
|
||||
You can create a basic demo site by running this command:
|
||||
@@ -166,18 +141,6 @@ You first need to create a superuser account:
|
||||
$ make superuser
|
||||
```
|
||||
|
||||
## Documentation 📚
|
||||
|
||||
Additional documentation is available in the `docs/` directory:
|
||||
|
||||
- [LLM Configuration](docs/llm-configuration.md) - Configure Large Language Models and providers
|
||||
- [Attachments](docs/attachments.md) - How to use attachments in conversations
|
||||
- [Tools for Agents](docs/tools.md) - Available tools and how to add new ones
|
||||
- [Environment Variables](docs/env.md) - All available environment variables
|
||||
- [Installation Guide](docs/installation.md) - Deploy on a Kubernetes cluster
|
||||
- [Theming](docs/theming.md) - Customize the application appearance
|
||||
- [Architecture](docs/architecture.md) - Technical architecture overview
|
||||
|
||||
## Licence 📝
|
||||
|
||||
This work is released under the MIT License (see [LICENSE](https://github.com/suitenumerique/conversations/blob/main/LICENSE)).
|
||||
|
||||
@@ -11,22 +11,3 @@ services:
|
||||
image: conversations:frontend-production
|
||||
ports:
|
||||
- "3000:3000"
|
||||
|
||||
openmockllm-mistral:
|
||||
user: "${DOCKER_USER:-1000}"
|
||||
build:
|
||||
context: .
|
||||
dockerfile: ./src/OpenMockLLM/Dockerfile
|
||||
image: conversations:openmockllm-mistral
|
||||
command:
|
||||
- openmockllm
|
||||
- --host
|
||||
- "0.0.0.0"
|
||||
- --port
|
||||
- "8000"
|
||||
- --backend
|
||||
- mistral
|
||||
- --model-name
|
||||
- mistral-mock
|
||||
ports:
|
||||
- "8900:8000"
|
||||
|
||||
-34
@@ -21,38 +21,6 @@ services:
|
||||
ports:
|
||||
- "1081:1080"
|
||||
|
||||
minio:
|
||||
user: ${DOCKER_USER:-1000}
|
||||
image: minio/minio
|
||||
environment:
|
||||
- MINIO_ROOT_USER=conversations
|
||||
- MINIO_ROOT_PASSWORD=password
|
||||
ports:
|
||||
- '9000:9000'
|
||||
- '9001:9001'
|
||||
healthcheck:
|
||||
test: ["CMD", "mc", "ready", "local"]
|
||||
interval: 1s
|
||||
timeout: 20s
|
||||
retries: 300
|
||||
entrypoint: ""
|
||||
command: minio server --console-address :9001 /data
|
||||
volumes:
|
||||
- ./data/media:/data
|
||||
|
||||
createbuckets:
|
||||
image: minio/mc
|
||||
depends_on:
|
||||
minio:
|
||||
condition: service_healthy
|
||||
restart: true
|
||||
entrypoint: >
|
||||
sh -c "
|
||||
/usr/bin/mc alias set conversations http://minio:9000 conversations password && \
|
||||
/usr/bin/mc mb conversations/conversations-media-storage && \
|
||||
/usr/bin/mc version enable conversations/conversations-media-storage && \
|
||||
exit 0;"
|
||||
|
||||
app-dev:
|
||||
build:
|
||||
context: .
|
||||
@@ -82,8 +50,6 @@ services:
|
||||
condition: service_started
|
||||
redis:
|
||||
condition: service_started
|
||||
createbuckets:
|
||||
condition: service_started
|
||||
|
||||
nginx:
|
||||
image: nginx:1.25
|
||||
|
||||
@@ -4,39 +4,6 @@ server {
|
||||
server_name localhost;
|
||||
charset utf-8;
|
||||
|
||||
# Proxy auth for media
|
||||
location /media/ {
|
||||
# Auth request configuration
|
||||
auth_request /media-auth;
|
||||
auth_request_set $authHeader $upstream_http_authorization;
|
||||
auth_request_set $authDate $upstream_http_x_amz_date;
|
||||
auth_request_set $authContentSha256 $upstream_http_x_amz_content_sha256;
|
||||
|
||||
# Pass specific headers from the auth response
|
||||
proxy_set_header Authorization $authHeader;
|
||||
proxy_set_header X-Amz-Date $authDate;
|
||||
proxy_set_header X-Amz-Content-SHA256 $authContentSha256;
|
||||
|
||||
# Get resource from Minio
|
||||
proxy_pass http://minio:9000/conversations-media-storage/;
|
||||
proxy_set_header Host minio:9000;
|
||||
|
||||
add_header Content-Security-Policy "default-src 'none'" always;
|
||||
}
|
||||
|
||||
location /media-auth {
|
||||
proxy_pass http://app-dev:8000/api/v1.0/chats/media-auth/;
|
||||
proxy_set_header Host $host;
|
||||
proxy_set_header X-Real-IP $remote_addr;
|
||||
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
|
||||
proxy_set_header X-Original-URL $request_uri;
|
||||
|
||||
# Prevent the body from being passed
|
||||
proxy_pass_request_body off;
|
||||
proxy_set_header Content-Length "";
|
||||
proxy_set_header X-Original-Method $request_method;
|
||||
}
|
||||
|
||||
location / {
|
||||
proxy_pass http://keycloak:8080;
|
||||
proxy_set_header Host $host;
|
||||
|
||||
@@ -7,8 +7,7 @@ flowchart TD
|
||||
User -- HTTP --> Front("Frontend (NextJS SPA)")
|
||||
Front -- REST API --> Back("Backend (Django)")
|
||||
Front -- OIDC --> Back -- OIDC ---> OIDC("Keycloak / ProConnect")
|
||||
Back -- REST API --> Yserver
|
||||
Back --> DB("Database (PostgreSQL)")
|
||||
Back --> Cache("Cache (Redis)")
|
||||
Back ----> S3("Minio (S3)")
|
||||
Back -- REST API --> LLM("LLM Providers")
|
||||
Back <--> Celery --> DB
|
||||
```
|
||||
|
||||
@@ -1,400 +0,0 @@
|
||||
# Conversation Attachments
|
||||
|
||||
This document describes how conversation attachments work in the Conversations application, including the upload process, security measures, and how documents are processed for use with Large Language Models (LLMs).
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Supported Attachment Types](#supported-attachment-types)
|
||||
- [Architecture & Flow](#architecture--flow)
|
||||
- [High-Level Overview](#high-level-overview)
|
||||
- [Detailed Technical Flow](#detailed-technical-flow)
|
||||
- [Security & Validation](#security--validation)
|
||||
- [MIME Type Validation](#mime-type-validation)
|
||||
- [Malware Detection](#malware-detection)
|
||||
- [Document Processing for LLMs](#document-processing-for-llms)
|
||||
- [Image Attachments](#image-attachments)
|
||||
- [PDF Documents](#pdf-documents)
|
||||
- [Other Document Types](#other-document-types)
|
||||
- [Configuration](#configuration)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
Conversations allows users to attach files to their conversations with the AI assistant. These attachments can be:
|
||||
- **Images** (displayed directly to vision-capable LLMs)
|
||||
- **PDF documents** (sent as document URLs to the LLM)
|
||||
- **Other documents** (converted to text and indexed for semantic search)
|
||||
|
||||
The attachment system uses **S3-compatible object storage** (such as MinIO in development) to store files securely.
|
||||
The backend generates **presigned URLs** that allow the frontend to upload files directly to the storage,
|
||||
without routing the file data through the backend server.
|
||||
|
||||
Note about documents: The system uses a tool called **MarkItDown** to convert various document formats
|
||||
(Word, Excel, PowerPoint, text files, etc.) into Markdown text for processing by LLMs. When at least
|
||||
one non-PDF/image document is attached, the system enables:
|
||||
- a **Retrieval-Augmented Generation (RAG)** search tool to allow the LLM to query relevant sections of the documents.
|
||||
- a **summarization tool** to provide document summaries on user request.
|
||||
⚠️ naive implementation at the moment, needs improvement before being used in production.
|
||||
|
||||
## Supported Attachment Types
|
||||
The following attachment types are supported:
|
||||
- **Images**: `image/png`, `image/jpeg`, `image/gif`, `image/webp`.
|
||||
- **PDF documents**: `application/pdf`
|
||||
- **Other documents**:
|
||||
- Microsoft Word: `application/vnd.openxmlformats-officedocument.wordprocessingml.document`
|
||||
- Microsoft Excel: `application/vnd.openxmlformats-officedocument.spreadsheetml.sheet`
|
||||
- Microsoft PowerPoint: `application/vnd.openxmlformats-officedocument.presentationml.presentation`
|
||||
- Text files: `text/plain`, `text/markdown`, `text/csv`
|
||||
|
||||
**Warning**: The current implementation for PDF expects the LLM to be able to manage them. We need to
|
||||
improve the handling of PDFs in case the LLM cannot process them natively.
|
||||
|
||||
**Todo**:
|
||||
- Add support for more file types and improve document processing workflows.
|
||||
- Allow PDF management via RAG search when the LLM cannot handle them natively.
|
||||
- Allow file type restrictions based on model settings, instead of globally.
|
||||
- Improve the summarization tool to provide better summaries and handle larger documents.
|
||||
- Start file upload right away when the user selects a file, instead of waiting for the user to send the message.
|
||||
|
||||
|
||||
---
|
||||
|
||||
## Architecture & Flow
|
||||
|
||||
### High-Level Overview
|
||||
|
||||
```
|
||||
┌─────────────┐ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐
|
||||
│ Frontend │ │ Backend │ │ S3 Storage │ │ Malware Det.│
|
||||
└──────┬──────┘ └──────┬──────┘ └──────┬──────┘ └──────┬──────┘
|
||||
│ │ │ │
|
||||
│ 1. Create attachment│ │ │
|
||||
├────────────────────>│ │ │
|
||||
│ │ │ │
|
||||
│ 2. Return presigned │ │ │
|
||||
│ URL for upload │ │ │
|
||||
│<────────────────────┤ │ │
|
||||
│ │ │ │
|
||||
│ 3. Upload file │ │ │
|
||||
│ directly to S3 │ │ │
|
||||
├──────────────────────────────────────────>│ │
|
||||
│ │ │ │
|
||||
│ 4. Notify upload │ │ │
|
||||
│ completed │ │ │
|
||||
├────────────────────>│ │ │
|
||||
│ │ │ │
|
||||
│ │ 5. Detect MIME type │ │
|
||||
│ ├────────────────────>│ │
|
||||
│ │ │ │
|
||||
│ │ 6. Scan for malware │ │
|
||||
│ ├──────────────────────────────────────────>│
|
||||
│ │ │ │
|
||||
│ │ 7. Update status │ │
|
||||
│ 8. Return status │<──────────────────────────────────────────┤
|
||||
│<────────────────────┤ │ │
|
||||
│ │ │ │
|
||||
```
|
||||
|
||||
### Detailed Technical Flow
|
||||
|
||||
#### Step 1: Attachment Creation Request
|
||||
|
||||
When a user selects a file to upload, the frontend sends a POST request to create an attachment record:
|
||||
|
||||
**Endpoint**: `POST /api/conversations/{conversation_id}/attachments/`
|
||||
|
||||
**Request payload**:
|
||||
```json
|
||||
{
|
||||
"file_name": "document.pdf",
|
||||
"size": 1048576,
|
||||
"content_type": "application/pdf"
|
||||
}
|
||||
```
|
||||
|
||||
**Backend processing** (`ChatConversationAttachmentViewSet.perform_create`):
|
||||
1. Verifies the user owns the conversation
|
||||
2. Generates a unique UUID for the file
|
||||
3. Creates a storage key: `{conversation_id}/attachments/{uuid}.{extension}`
|
||||
4. Creates a database record with status `PENDING`
|
||||
|
||||
**Response**:
|
||||
```json
|
||||
{
|
||||
"id": "uuid-of-attachment",
|
||||
"key": "conversation-id/attachments/file-id.pdf",
|
||||
"file_name": "document.pdf",
|
||||
"size": 1048576,
|
||||
"upload_state": "pending",
|
||||
"policy": "https://s3.example.com/bucket/...?presigned-params"
|
||||
}
|
||||
```
|
||||
|
||||
The `policy` field contains a **presigned URL** valid for a limited time (configured by `AWS_S3_UPLOAD_POLICY_EXPIRATION`).
|
||||
|
||||
#### Step 2: Direct Upload to S3
|
||||
|
||||
The frontend uses the presigned URL to upload the file directly to S3 storage using a PUT request.
|
||||
|
||||
**Technical details**:
|
||||
- The presigned URL includes authentication parameters
|
||||
- The upload is done with `Content-Type` header matching the file's MIME type
|
||||
- No backend involvement in the data transfer
|
||||
|
||||
#### Step 3: Upload Completion Notification
|
||||
|
||||
After successful upload, the frontend notifies the backend:
|
||||
|
||||
**Endpoint**: `POST /api/conversations/{conversation_id}/attachments/{attachment_id}/upload-ended/`
|
||||
|
||||
**Backend processing** (`ChatConversationAttachmentViewSet.upload_ended`):
|
||||
|
||||
1. **MIME Type Detection** (`chat/views.py`):
|
||||
```python
|
||||
mime_detector = magic.Magic(mime=True)
|
||||
with default_storage.open(attachment.key, "rb") as file:
|
||||
mimetype = mime_detector.from_buffer(file.read(2048))
|
||||
size = file.size
|
||||
```
|
||||
|
||||
Uses `python-magic` to detect the actual MIME type from file content (first 2048 bytes).
|
||||
|
||||
2. **Update attachment status**:
|
||||
- Status: `PENDING` → `ANALYZING`
|
||||
- Store detected MIME type and actual file size
|
||||
|
||||
3. **Trigger Malware Detection**:
|
||||
```python
|
||||
malware_detection.analyse_file(
|
||||
attachment.key,
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
```
|
||||
|
||||
#### Step 4: Malware Detection Callbacks
|
||||
|
||||
The malware detection service (configurable via `MALWARE_DETECTION_BACKEND`) scans the file and calls one of three callbacks:
|
||||
|
||||
**Safe file** (`conversation_safe_attachment_callback`):
|
||||
- Status: `ANALYZING` → `READY`
|
||||
- File is ready for use
|
||||
|
||||
**Unsafe file** (`conversation_unsafe_attachment_callback`):
|
||||
- Status: `ANALYZING` → `SUSPICIOUS`
|
||||
- File is quarantined and not accessible
|
||||
- Security log entry created
|
||||
|
||||
**Unknown status** (`unknown_attachment_callback`):
|
||||
- Handles special cases (e.g., file too large to analyze)
|
||||
- Status: `ANALYZING` → `FILE_TOO_LARGE_TO_ANALYZE`
|
||||
|
||||
---
|
||||
|
||||
## Security & Validation
|
||||
|
||||
For now, the system is not intended to host user-uploaded files for public download.
|
||||
All files are stored in private S3 buckets with presigned URLs for controlled access and only
|
||||
the owner of the conversation/the uploader can access them, so the risk is quite low around bad use of
|
||||
the attachment system.
|
||||
|
||||
Also, the document content is sent to the LLM and does not prevent any prompt injection attacks, which is not
|
||||
an issue specific to the attachment system but to the overall design of LLM-based applications and should be
|
||||
addressed globally. Also for the moment, the system does not have any action tools that could be used to execute
|
||||
malicious code based on document content.
|
||||
|
||||
### Malware Detection
|
||||
|
||||
The malware detection system is **pluggable** and configurable, allowing different backends to be used.
|
||||
By default, a `DummyBackend` is provided that marks all files as safe.
|
||||
|
||||
⚠️ The current implementation does not disallow any file types or status from being used in conversations.
|
||||
This is a potential security risk and should be addressed in future versions.
|
||||
|
||||
---
|
||||
|
||||
## Document Processing for LLMs
|
||||
|
||||
When a user sends a message with attachments, the system processes them differently based on their type:
|
||||
|
||||
### Image Attachments
|
||||
|
||||
**MIME types**: `image/png`, `image/jpeg`, `image/gif`, `image/webp`, etc.
|
||||
|
||||
**Processing flow**:
|
||||
|
||||
1. **URL Conversion**: Local media URLs are converted to presigned S3 URLs before sending to the LLM:
|
||||
```python
|
||||
# From: chat/agents/local_media_url_processors.py
|
||||
content.url = generate_retrieve_policy(key)
|
||||
```
|
||||
|
||||
2. **Sent to LLM**: Images are sent as `ImageUrl` objects in the prompt:
|
||||
```python
|
||||
ImageUrl(
|
||||
url="https://s3.example.com/bucket/key?presigned-params",
|
||||
identifier="file-id.png",
|
||||
)
|
||||
```
|
||||
|
||||
3. **Vision models** can analyze the image content directly.
|
||||
|
||||
4. **Response processing**: After the LLM responds, presigned URLs are converted back to local URLs for storage:
|
||||
```python
|
||||
# Mapping: presigned_url -> /media-key/{conversation_id}/attachments/{file_id}.png
|
||||
```
|
||||
|
||||
### PDF Documents
|
||||
|
||||
**MIME type**: `application/pdf`
|
||||
|
||||
**Processing flow**:
|
||||
|
||||
1. **Direct URL passing**: PDFs are sent as `DocumentUrl` objects :
|
||||
```python
|
||||
DocumentUrl(
|
||||
url="https://s3.example.com/bucket/key?presigned-params",
|
||||
identifier="file-id.pdf",
|
||||
)
|
||||
```
|
||||
|
||||
2. **LLM processing**: Compatible LLMs can:
|
||||
- Extract and read text from PDFs
|
||||
- Understand document structure
|
||||
- Answer questions about the content
|
||||
|
||||
3. **No conversion needed**: PDFs are passed directly without preprocessing.
|
||||
|
||||
### Other Document Types
|
||||
|
||||
**MIME types**: Word documents, Excel spreadsheets, PowerPoint, text files, Markdown, etc.
|
||||
|
||||
**Processing flow**:
|
||||
|
||||
1. **Document parsing**: When a document is uploaded, it's parsed using the `AlbertRagBackend` class.
|
||||
|
||||
2. **Conversion to Markdown**: Documents are converted using **MarkItDown** library or using the "Albert API" for PDFs.
|
||||
|
||||
3. **RAG (Retrieval-Augmented Generation)**:
|
||||
- Converted text is indexed in a vector database
|
||||
- The LLM uses a `document_rag_search` tool to query relevant sections
|
||||
- Only relevant chunks are sent to the LLM to fit context windows
|
||||
|
||||
4. **Summarization tool** if needed.
|
||||
|
||||
### Processing Strategy Decision Tree
|
||||
|
||||
**Decision logic**:
|
||||
- **No documents**: Standard conversation
|
||||
- **Images**: Send as direct (presigned) URLs to the LLM
|
||||
- **Only PDFs**: Send as direct (presigned) URLs to the LLM
|
||||
- **Other documents present**: Enable RAG search tool + convert to Markdown
|
||||
|
||||
---
|
||||
|
||||
## Configuration
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------------------------------------------|----------------|------------------------------------------------------------|
|
||||
| `ATTACHMENT_MAX_SIZE` | Configurable | Maximum file size in bytes |
|
||||
| `ATTACHMENT_CHECK_UNSAFE_MIME_TYPES_ENABLED` | `True` | Enable/disable MIME type validation |
|
||||
| `AWS_S3_UPLOAD_POLICY_EXPIRATION` | 3600 | Presigned URL expiration (seconds) |
|
||||
| `AWS_S3_RETRIEVE_POLICY_EXPIRATION` | 3600 | Presigned retrieval URL expiration (seconds) |
|
||||
| `AWS_S3_DOMAIN_REPLACE` | None | Alternative S3 domain for presigned URLs (for development) |
|
||||
| `MALWARE_DETECTION_BACKEND` | `DummyBackend` | Malware scanning backend class |
|
||||
| `MALWARE_DETECTION_PARAMETERS` | `{}` | Backend-specific configuration |
|
||||
| `RAG_FILES_ACCEPTED_FORMATS` | See below | List of MIME types accepted for file uploads |
|
||||
|
||||
#### RAG_FILES_ACCEPTED_FORMATS
|
||||
|
||||
This environment variable controls which file types users are allowed to upload as attachments to conversations.
|
||||
|
||||
**Configuration**:
|
||||
- **Type**: List of strings (comma-separated MIME types when using environment variable)
|
||||
- **Default value**: Includes a comprehensive list of document and image formats:
|
||||
- Microsoft Office documents (`.docx`, `.pptx`, `.xlsx`, `.xls`)
|
||||
- Text files (`.txt`, `.csv`)
|
||||
- PDF documents (`.pdf`)
|
||||
- HTML files
|
||||
- Markdown files (`.md`)
|
||||
- Outlook messages (`.msg`)
|
||||
- Images (`.jpeg`, `.png`, `.gif`, `.webp`)
|
||||
|
||||
**Example configuration**:
|
||||
```ini
|
||||
# In environment variable (comma-separated)
|
||||
RAG_FILES_ACCEPTED_FORMATS="application/pdf,text/plain,image/png,image/jpeg"
|
||||
```
|
||||
|
||||
```python
|
||||
# In Django settings (as a Python list)
|
||||
RAG_FILES_ACCEPTED_FORMATS = [
|
||||
"application/pdf",
|
||||
"text/plain",
|
||||
"image/png",
|
||||
"image/jpeg",
|
||||
]
|
||||
```
|
||||
|
||||
**How it's used**:
|
||||
1. **Backend**: The list is exposed via the `/api/v1.0/config/` endpoint as `chat_upload_accept` (MIME types joined with commas)
|
||||
2. **Frontend**: The configuration is used to validate files before upload in the chat interface:
|
||||
- Checks exact MIME type matches
|
||||
- Supports wildcard patterns (e.g., `image/*` for all image types)
|
||||
- Supports file extension patterns (e.g., `.pdf`)
|
||||
3. **User experience**: Files that don't match the accepted formats are rejected with a user-friendly error message
|
||||
|
||||
**Notes**:
|
||||
|
||||
- This setting controls frontend validation only. Backend validation should also be implemented for security.
|
||||
- Future improvements may include per-model file type restrictions.
|
||||
|
||||
### Storage Configuration
|
||||
|
||||
**MinIO (Development)**:
|
||||
```yaml
|
||||
# docker-compose.yml
|
||||
minio:
|
||||
image: minio/minio
|
||||
environment:
|
||||
MINIO_ROOT_USER: minioadmin
|
||||
MINIO_ROOT_PASSWORD: minioadmin
|
||||
command: server /data --console-address ":9001"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### LLM Cannot Access Image/PDF
|
||||
|
||||
**Possible causes**:
|
||||
- Presigned URL has expired
|
||||
- S3 storage is not accessible from the LLM provider
|
||||
- CORS configuration issues
|
||||
|
||||
**Solution**: Check `AWS_S3_RETRIEVE_POLICY_EXPIRATION` and S3 access policies.
|
||||
|
||||
### Document Not Appearing in RAG Search
|
||||
|
||||
**Possible causes**:
|
||||
- Document conversion failed
|
||||
- Vector database indexing failed
|
||||
|
||||
**Check logs**: Look for errors in `DocumentConverter` and RAG backend logs.
|
||||
|
||||
---
|
||||
|
||||
## Related Documentation
|
||||
|
||||
- [Installation Guide](installation.md) - S3 storage setup
|
||||
- [LLM Configuration](llm-configuration.md) - Model capabilities for attachments
|
||||
- [Architecture](architecture.md) - System overview
|
||||
- [Tools](tools.md) - Document search and RAG tools
|
||||
|
||||
+9
-14
@@ -10,6 +10,7 @@ These are the environment variables you can set for the `conversations-backend`
|
||||
|-------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------|
|
||||
| DJANGO_ALLOWED_HOSTS | allowed hosts | [] |
|
||||
| DJANGO_SECRET_KEY | secret key | |
|
||||
| DJANGO_SERVER_TO_SERVER_API_TOKENS | | [] |
|
||||
| DB_ENGINE | engine to use for database connections | django.db.backends.postgresql_psycopg2 |
|
||||
| DB_NAME | name of the database | conversations |
|
||||
| DB_USER | user to authenticate with | dinum |
|
||||
@@ -18,16 +19,12 @@ These are the environment variables you can set for the `conversations-backend`
|
||||
| DB_PORT | port of the database | 5432 |
|
||||
| MEDIA_BASE_URL | | |
|
||||
| STORAGES_STATICFILES_BACKEND | | whitenoise.storage.CompressedManifestStaticFilesStorage |
|
||||
| AWS_S3_ENDPOINT_URL | S3 endpoint | |
|
||||
| AWS_S3_ACCESS_KEY_ID | access id for s3 endpoint | |
|
||||
| AWS_S3_SECRET_ACCESS_KEY | access key for s3 endpoint | |
|
||||
| AWS_S3_REGION_NAME | region name for s3 endpoint | |
|
||||
| AWS_STORAGE_BUCKET_NAME | bucket name for s3 endpoint | conversations-media-storage |
|
||||
| ATTACHMENT_MAX_SIZE | maximum size of document in bytes | 10485760 |
|
||||
| DOCUMENT_IMAGE_MAX_SIZE | maximum size of document in bytes | 10485760 |
|
||||
| LANGUAGE_CODE | default language | en-us |
|
||||
| API_USERS_LIST_THROTTLE_RATE_SUSTAINED | throttle rate for api | 180/hour |
|
||||
| API_USERS_LIST_THROTTLE_RATE_BURST | throttle rate for api on burst | 30/minute |
|
||||
| SPECTACULAR_SETTINGS_ENABLE_DJANGO_DEPLOY_CHECK | | false |
|
||||
| TRASHBIN_CUTOFF_DAYS | trashbin cutoff | 30 |
|
||||
| DJANGO_EMAIL_BACKEND | email backend library | django.core.mail.backends.smtp.EmailBackend |
|
||||
| DJANGO_EMAIL_BRAND_NAME | brand name for email | |
|
||||
| DJANGO_EMAIL_HOST | host name of email | |
|
||||
@@ -74,14 +71,12 @@ These are the environment variables you can set for the `conversations-backend`
|
||||
| OIDC_USERINFO_FULLNAME_FIELDS | OIDC token claims to create full name | ["first_name", "last_name"] |
|
||||
| OIDC_USERINFO_SHORTNAME_FIELD | OIDC token claims to create shortname | first_name |
|
||||
| ALLOW_LOGOUT_GET_METHOD | Allow get logout method | true |
|
||||
| LLM_CONFIGURATION_FILE_PATH | Path to the LLM configuration JSON file. See [LLM Configuration](llm-configuration.md) for details | <BASE_DIR>/conversations/configuration/llm/default.json |
|
||||
| LLM_DEFAULT_MODEL_HRID | HRID of the model used for conversations | default-model |
|
||||
| LLM_SUMMARIZATION_MODEL_HRID | HRID of the model used for summarization | default-summarization-model |
|
||||
| AI_API_KEY | AI API key to be used for the default provider (used in default LLM configuration, not for production use) | |
|
||||
| AI_BASE_URL | OpenAI compatible AI base URL (used in default LLM configuration, not for production use) | |
|
||||
| AI_MODEL | AI Model name to use (used in default LLM configuration, not for production use) | |
|
||||
| AI_AGENT_INSTRUCTIONS | Base instruction for the AI agent (used in default LLM configuration, not for production use) | You are a helpful assistant. Wrap formulas... |
|
||||
| AI_AGENT_TOOLS | List of enabled tools for the agent (used in default LLM configuration, not for production use) | [] |
|
||||
| AI_API_KEY | AI key to be used for AI Base url | |
|
||||
| AI_BASE_URL | OpenAI compatible AI base url | |
|
||||
| AI_MODEL | AI Model to use | |
|
||||
| AI_AGENT_INSTRUCTION | Base instruction for the AI agent | You are a helpful assistant |
|
||||
| Y_PROVIDER_API_KEY | Y provider API key | |
|
||||
| Y_PROVIDER_API_BASE_URL | Y Provider url | |
|
||||
| CONVERSION_API_ENDPOINT | Conversion API endpoint | convert-markdown |
|
||||
| CONVERSION_API_CONTENT_FIELD | Conversion api content field | content |
|
||||
| CONVERSION_API_TIMEOUT | Conversion api timeout | 30 |
|
||||
|
||||
@@ -9,6 +9,7 @@ backend:
|
||||
DJANGO_CSRF_TRUSTED_ORIGINS: https://conversations.127.0.0.1.nip.io
|
||||
DJANGO_CONFIGURATION: Feature
|
||||
DJANGO_ALLOWED_HOSTS: conversations.127.0.0.1.nip.io
|
||||
DJANGO_SERVER_TO_SERVER_API_TOKENS: secret-api-key
|
||||
DJANGO_SECRET_KEY: AgoodOrAbadKey
|
||||
DJANGO_SETTINGS_MODULE: conversations.settings
|
||||
DJANGO_SUPERUSER_PASSWORD: admin
|
||||
@@ -47,10 +48,6 @@ backend:
|
||||
POSTGRES_USER: dinum
|
||||
POSTGRES_PASSWORD: pass
|
||||
REDIS_URL: redis://default:pass@redis-master:6379/1
|
||||
AWS_S3_ENDPOINT_URL: http://minio.conversations.svc.cluster.local:9000
|
||||
AWS_S3_ACCESS_KEY_ID: root
|
||||
AWS_S3_SECRET_ACCESS_KEY: password
|
||||
AWS_STORAGE_BUCKET_NAME: conversations-media-storage
|
||||
STORAGES_STATICFILES_BACKEND: django.contrib.staticfiles.storage.StaticFilesStorage
|
||||
|
||||
migrate:
|
||||
@@ -116,18 +113,3 @@ ingress:
|
||||
ingressAdmin:
|
||||
enabled: true
|
||||
host: conversations.127.0.0.1.nip.io
|
||||
|
||||
ingressMedia:
|
||||
enabled: true
|
||||
host: conversations.127.0.0.1.nip.io
|
||||
|
||||
annotations:
|
||||
nginx.ingress.kubernetes.io/auth-url: https://conversations.127.0.0.1.nip.io/api/v1.0/chats/media-auth/
|
||||
nginx.ingress.kubernetes.io/auth-response-headers: "Authorization, X-Amz-Date, X-Amz-Content-SHA256"
|
||||
nginx.ingress.kubernetes.io/upstream-vhost: minio.conversations.svc.cluster.local:9000
|
||||
nginx.ingress.kubernetes.io/rewrite-target: /conversations-media-storage/$1
|
||||
|
||||
serviceMedia:
|
||||
host: minio.conversations.svc.cluster.local
|
||||
port: 9000
|
||||
|
||||
|
||||
+1
-21
@@ -7,8 +7,7 @@ This document is a step-by-step guide that describes how to install Conversation
|
||||
- k8s cluster with an nginx-ingress controller
|
||||
- an OIDC provider (if you don't have one, we provide an example)
|
||||
- a PostgreSQL server (if you don't have one, we provide an example)
|
||||
- a Redis server (if you don't have one, we provide an example)
|
||||
- a S3 bucket (if you don't have one, we provide an example)
|
||||
- a Memcached server (if you don't have one, we provide an example)
|
||||
|
||||
### Test cluster
|
||||
|
||||
@@ -173,22 +172,6 @@ POSTGRES_USER: dinum
|
||||
POSTGRES_PASSWORD: pass
|
||||
```
|
||||
|
||||
### Find s3 bucket connection values
|
||||
|
||||
Conversations uses an s3 bucket to store documents, so if you have a provider obtain the necessary information to use it. If you don't, you can install a local minio testing environment as follow:
|
||||
|
||||
```
|
||||
$ helm install minio oci://registry-1.docker.io/bitnamicharts/minio -f examples/minio.values.yaml
|
||||
$ kubectl get po
|
||||
NAME READY STATUS RESTARTS AGE
|
||||
keycloak-0 1/1 Running 0 38m
|
||||
keycloak-postgresql-0 1/1 Running 0 38m
|
||||
minio-84f5c66895-bbhsk 1/1 Running 0 42s
|
||||
minio-provisioning-2b5sq 0/1 Completed 0 42s
|
||||
postgresql-0 1/1 Running 0 24m
|
||||
redis-master-0 1/1 Running 0 10m
|
||||
```
|
||||
|
||||
## Deployment
|
||||
|
||||
Now you are ready to deploy Conversations without AI. AI requires more dependencies (OpenAI API). To deploy Conversations you need to provide all previous information to the helm chart.
|
||||
@@ -205,8 +188,6 @@ conversations-conversations-backend-migrate-c949s 0/1 Completed
|
||||
conversations-conversations-frontend-6749f644f7-p5s42 1/1 Running 0 79s
|
||||
keycloak-0 1/1 Running 0 48m
|
||||
keycloak-postgresql-0 1/1 Running 0 48m
|
||||
minio-84f5c66895-bbhsk 1/1 Running 0 10m
|
||||
minio-provisioning-2b5sq 0/1 Completed 0 10m
|
||||
postgresql-0 1/1 Running 0 34m
|
||||
redis-master-0 1/1 Running 0 20m
|
||||
```
|
||||
@@ -220,7 +201,6 @@ $ kubectl get ingress
|
||||
NAME CLASS HOSTS ADDRESS PORTS AGE
|
||||
conversations-conversations <none> conversations.127.0.0.1.nip.io localhost 80, 443 114s
|
||||
conversations-conversations-admin <none> conversations.127.0.0.1.nip.io localhost 80, 443 114s
|
||||
conversations-conversations-media <none> conversations.127.0.0.1.nip.io localhost 80, 443 114s
|
||||
conversations-conversations-ws <none> conversations.127.0.0.1.nip.io localhost 80, 443 114s
|
||||
keycloak <none> keycloak.127.0.0.1.nip.io localhost 80 49m
|
||||
```
|
||||
|
||||
@@ -1,412 +0,0 @@
|
||||
# LLM Configuration
|
||||
|
||||
This document describes how to configure Large Language Models (LLMs) in Conversations via the configuration file.
|
||||
|
||||
## Overview
|
||||
|
||||
Conversations uses a JSON configuration file to define LLM models and providers. This approach allows you to:
|
||||
- Configure multiple LLM models from different providers
|
||||
- Switch between models without code changes
|
||||
- Customize model-specific settings like temperature, max tokens, and system prompts
|
||||
- Enable or disable models dynamically
|
||||
|
||||
The overall structure consists of two main sections: `providers` and `models`.
|
||||
Settings for models, provides customization through `settings` and `profile`, which corresponds to the
|
||||
Pydantic AI model settings and profile. While we currently not use those settings extensively,
|
||||
they are available for future use and advanced configurations, please reach us if you face any problem using them.
|
||||
|
||||
## Configuration File Location
|
||||
|
||||
The default LLM configuration file is located at:
|
||||
```
|
||||
src/backend/conversations/configuration/llm/default.json
|
||||
```
|
||||
|
||||
You can override this location by setting the `LLM_CONFIGURATION_FILE_PATH` environment variable, but be careful as
|
||||
this path must be accessible by the backend application _inside the docker image_:
|
||||
``` ini
|
||||
LLM_CONFIGURATION_FILE_PATH=/path/to/your/llm/config.json
|
||||
```
|
||||
|
||||
## Default Behavior
|
||||
|
||||
### Default Configuration
|
||||
|
||||
The default configuration file is useful for local development and running the test, while it can be used
|
||||
in production, we suggest to create a specific one for production and replace the `settings.` values with
|
||||
`environ.` one.
|
||||
|
||||
The default configuration file (`default.json`) includes:
|
||||
|
||||
1. **Two default models**:
|
||||
- `default-model`: The primary conversational model used for chat interactions
|
||||
- `default-summarization-model`: A specialized model for summarizing conversations
|
||||
|
||||
2. **One default provider**:
|
||||
- `default-provider`: An OpenAI-compatible provider that uses environment variables for configuration
|
||||
|
||||
### Environment Variable Integration
|
||||
|
||||
The configuration uses dynamic value resolution with two special prefixes:
|
||||
|
||||
- `settings.VARIABLE_NAME`: Resolves to a Django setting value
|
||||
- `environ.VARIABLE_NAME`: Resolves to an environment variable value
|
||||
|
||||
For example, in the default configuration:
|
||||
```json
|
||||
{
|
||||
"model_name": "settings.AI_MODEL",
|
||||
"system_prompt": "settings.AI_AGENT_INSTRUCTIONS",
|
||||
"tools": "settings.AI_AGENT_TOOLS"
|
||||
}
|
||||
```
|
||||
|
||||
This allows to configure models in tests using the setting override mechanism from Django/Pytest (but might be replaced
|
||||
later with a simple override of the full configuration like it's done in some tests already).
|
||||
|
||||
### Required Environment Variables
|
||||
|
||||
For the default configuration to work, you need to set these environment variables:
|
||||
|
||||
| Variable | Description | Example |
|
||||
|-------------------------------|----------------------------------------|-----------------------------|
|
||||
| `AI_API_KEY` | API key for the default provider | `sk-...` |
|
||||
| `AI_BASE_URL` | Base URL for the OpenAI-compatible API | `https://api.openai.com/v1` |
|
||||
| `AI_MODEL` | Model name to use | `gpt-4o-mini` |
|
||||
|
||||
### Optional Environment Variables
|
||||
|
||||
If you want to customize the agent behavior and tools, you can set these optional environment variables
|
||||
(defaults are provided in the default configuration):
|
||||
|
||||
| Variable | Description | Default |
|
||||
|-------------------------------|----------------------------------------|-------------------|
|
||||
| `AI_AGENT_INSTRUCTIONS` | System prompt for the agent | see `settings.py` |
|
||||
| `AI_AGENT_TOOLS` | List of enabled tools | `[]` |
|
||||
| `SUMMARIZATION_SYSTEM_PROMPT` | Base prompt of the summarization agent | see `settings.py` |
|
||||
|
||||
### Model Selection
|
||||
|
||||
You can configure which models are used for specific tasks via environment variables:
|
||||
|
||||
| Variable | Description | Default |
|
||||
|--------------------------------|------------------------------------------|-------------------------------|
|
||||
| `LLM_DEFAULT_MODEL_HRID` | HRID of the model used for conversations | `default-model` |
|
||||
| `LLM_SUMMARIZATION_MODEL_HRID` | HRID of the model used for summarization | `default-summarization-model` |
|
||||
|
||||
## Configuration Structure
|
||||
|
||||
The configuration file has two main sections:
|
||||
|
||||
### 1. Providers
|
||||
|
||||
Providers define the API endpoints and authentication for LLM services.
|
||||
|
||||
```json
|
||||
{
|
||||
"providers": [
|
||||
{
|
||||
"hrid": "unique-provider-id",
|
||||
"base_url": "https://api.example.com/v1",
|
||||
"api_key": "environ.API_KEY_VAR",
|
||||
"kind": "openai"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Provider Fields:**
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|------------|--------|----------|---------------------------------------------------------|
|
||||
| `hrid` | string | Yes | Unique identifier for the provider |
|
||||
| `base_url` | string | Yes | API base URL (can use `settings.` or `environ.` prefix) |
|
||||
| `api_key` | string | Yes | API authentication key (use `environ.` here) |
|
||||
| `kind` | string | Yes | Provider type: `openai` or `mistral` |
|
||||
|
||||
### 2. Models
|
||||
|
||||
Models define the LLMs available in your application.
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "unique-model-id",
|
||||
"model_name": "gpt-4o-mini",
|
||||
"human_readable_name": "GPT-4o Mini",
|
||||
"provider_name": "unique-provider-id",
|
||||
"profile": null,
|
||||
"settings": {},
|
||||
"is_active": true,
|
||||
"icon": null,
|
||||
"system_prompt": "You are a helpful assistant",
|
||||
"tools": []
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Model Fields:**
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|-----------------------|--------------|----------|-----------------------------------------------------------------------------------------------------|
|
||||
| `hrid` | string | Yes | Unique identifier for the model |
|
||||
| `model_name` | string | Yes | Name of the model as recognized by the provider (can use `settings.` or `environ.` prefix) |
|
||||
| `human_readable_name` | string | Yes | Display name shown to users |
|
||||
| `provider_name` | string | No* | Reference to a provider's `hrid` |
|
||||
| `provider` | object | No* | Inline provider definition (alternative to `provider_name`) |
|
||||
| `profile` | object | No | Model-specific capabilities and settings |
|
||||
| `settings` | object | No | Model inference settings (temperature, max_tokens, etc.) |
|
||||
| `is_active` | boolean | Yes | Whether the model is available for use |
|
||||
| `icon` | string/array | No | Base64-encoded icon or array of icon parts |
|
||||
| `system_prompt` | string | Yes | Default system prompt for the model (can use `settings.` or `environ.` prefix) |
|
||||
| `tools` | array | Yes | List of enabled tools for this model (can use `settings.` or `environ.` prefix for the whole array) |
|
||||
| `supports_streaming` | boolean | No | Whether the model supports streaming responses |
|
||||
|
||||
\* Either `provider_name` or `provider` must be set, unless `model_name` is in the format `<provider>:<model>`.
|
||||
|
||||
## Adding New Models
|
||||
|
||||
### Example 1: Adding a New OpenAI Model
|
||||
|
||||
To add a new OpenAI model using the existing default provider:
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
// ...existing models...
|
||||
{
|
||||
"hrid": "gpt-4-turbo",
|
||||
"model_name": "gpt-4-turbo-preview",
|
||||
"human_readable_name": "GPT-4 Turbo",
|
||||
"provider_name": "default-provider",
|
||||
"profile": null,
|
||||
"settings": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 4096
|
||||
},
|
||||
"is_active": true,
|
||||
"icon": null,
|
||||
"system_prompt": "You are an expert AI assistant.",
|
||||
"tools": ["web_search_brave_with_document_backend"],
|
||||
"supports_streaming": true
|
||||
}
|
||||
],
|
||||
"providers": [
|
||||
// ...existing providers...
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Example 2: Adding a Model using Pydantic AI format
|
||||
|
||||
To add a model with a specific provider using the default Pydantic AI format, you don't need to define the provider separately if you use the `model_name` format `<provider>:<model>`.
|
||||
|
||||
1. **Add the model without provider**:
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "claude-3-opus",
|
||||
"model_name": "anthropic:claude-3-opus-20240229",
|
||||
"human_readable_name": "Claude 3 Opus",
|
||||
"provider_name": null,
|
||||
"profile": null,
|
||||
"settings": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 4096
|
||||
},
|
||||
"is_active": true,
|
||||
"icon": null,
|
||||
"system_prompt": "You are Claude, a helpful AI assistant.",
|
||||
"tools": []
|
||||
}
|
||||
],
|
||||
"providers": []
|
||||
}
|
||||
```
|
||||
|
||||
2**Set the environment variable**:
|
||||
|
||||
Pydantic AI expects the API key in an environment variable named `ANTHROPIC_API_KEY` is this example, so set it accordingly:
|
||||
|
||||
```ini
|
||||
ANTHROPIC_API_KEY=your-api-key-here
|
||||
```
|
||||
|
||||
### Example 3: Adding a Mistral Model
|
||||
|
||||
For Mistral AI models using the Etalab platform:
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "mistral-large",
|
||||
"model_name": "mistral-large-latest",
|
||||
"human_readable_name": "Mistral Large (Etalab)",
|
||||
"provider_name": "mistral-etalab",
|
||||
"profile": null,
|
||||
"settings": {
|
||||
"temperature": 0.5,
|
||||
"max_tokens": 8192
|
||||
},
|
||||
"is_active": true,
|
||||
"icon": null,
|
||||
"system_prompt": "settings.AI_AGENT_INSTRUCTIONS",
|
||||
"tools": ["web_search_brave_with_document_backend"]
|
||||
}
|
||||
],
|
||||
"providers": [
|
||||
{
|
||||
"hrid": "mistral-etalab",
|
||||
"base_url": "https://api.mistral.etalab.gouv.fr/",
|
||||
"api_key": "environ.MISTRAL_ETALAB_API_KEY",
|
||||
"kind": "mistral"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Example 4: Using Inline Provider Definition
|
||||
|
||||
Instead of referencing a provider by name, you can define it inline if you use a unique configuration:
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "custom-model",
|
||||
"model_name": "custom-model-v1",
|
||||
"human_readable_name": "Custom Model",
|
||||
"provider": {
|
||||
"hrid": "custom-provider-inline",
|
||||
"base_url": "https://custom-api.example.com/v1",
|
||||
"api_key": "environ.CUSTOM_API_KEY",
|
||||
"kind": "openai"
|
||||
},
|
||||
"settings": {},
|
||||
"is_active": true,
|
||||
"icon": null,
|
||||
"system_prompt": "You are a custom assistant.",
|
||||
"tools": []
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
### Model Settings
|
||||
|
||||
The `settings` object supports various inference parameters:
|
||||
|
||||
```json
|
||||
{
|
||||
"settings": {
|
||||
"max_tokens": 4096,
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"timeout": 60.0,
|
||||
"parallel_tool_calls": true,
|
||||
"seed": 42,
|
||||
"presence_penalty": 0.0,
|
||||
"frequency_penalty": 0.0,
|
||||
"logit_bias": {},
|
||||
"stop_sequences": [],
|
||||
"extra_headers": {},
|
||||
"extra_body": {}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Model Profile
|
||||
|
||||
The `profile` object defines model capabilities:
|
||||
|
||||
```json
|
||||
{
|
||||
"profile": {
|
||||
"supports_tools": true,
|
||||
"supports_json_schema_output": true,
|
||||
"supports_json_object_output": true,
|
||||
"default_structured_output_mode": "json_schema",
|
||||
"thinking_tags": ["<thinking>", "</thinking>"],
|
||||
"ignore_streamed_leading_whitespace": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Available Tools
|
||||
|
||||
Tools can be specified in the `tools` array. Common tools include:
|
||||
- `web_search_brave_with_document_backend`: Web search using Brave API with document processing
|
||||
|
||||
You can also reference the tools list from Django settings:
|
||||
```json
|
||||
{
|
||||
"tools": "settings.AI_AGENT_TOOLS"
|
||||
}
|
||||
```
|
||||
|
||||
### Custom Icons
|
||||
|
||||
Icons can be provided as base64-encoded PNG images. For long strings, you can split them into an array:
|
||||
|
||||
```json
|
||||
{
|
||||
"icon": [
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABwAAAAcCAMAAABF0y+m",
|
||||
"AAAAn1BMVEUALosAKoovTZjw8vb////+9/jlPUniAAziABUAGIWbpsTwq7HhAAAA"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Validation
|
||||
|
||||
The configuration is validated when loaded. Common validation errors include:
|
||||
|
||||
- **Provider not found**: A model references a `provider_name` that doesn't exist in the `providers` array
|
||||
- **Missing provider**: Neither `provider_name` nor `provider` is specified, and `model_name` is not in `<provider>:<model>` format
|
||||
- **Environment variable not set**: A value using `environ.` prefix references an undefined environment variable
|
||||
- **Django setting not set**: A value using `settings.` prefix references an undefined Django setting
|
||||
- **Invalid provider kind**: The `kind` field must be either `openai` or `mistral`
|
||||
|
||||
## Testing Your Configuration
|
||||
|
||||
After modifying the configuration file, you can test it by:
|
||||
|
||||
1. **Checking for syntax errors**:
|
||||
```bash
|
||||
python -m json.tool src/backend/conversations/configuration/llm/default.json
|
||||
```
|
||||
|
||||
2. **Starting the application** and checking the logs for validation errors
|
||||
|
||||
3. **Using the Django shell** to load the configuration:
|
||||
```bash
|
||||
./bin/manage shell
|
||||
```
|
||||
```python
|
||||
from django.conf import settings
|
||||
models = settings.LLM_CONFIGURATIONS
|
||||
models.keys() # Should show all model HRIDs
|
||||
```
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Use environment variables** for sensitive data like API keys (with `environ.` prefix)
|
||||
2. **Use Django settings** for configurable values that may change between environments (with `settings.` prefix)
|
||||
3. **Keep provider definitions separate** from models to avoid duplication when using multiple models from the same provider
|
||||
4. **Set `is_active: false`** for models you want to keep in the configuration but temporarily disable
|
||||
5. **Use descriptive `hrid` values** that clearly identify the model and provider
|
||||
6. **Document custom configurations** in your deployment documentation
|
||||
7. **Test configuration changes** in a development environment before deploying to production
|
||||
|
||||
## See Also
|
||||
|
||||
- [Environment Variables Documentation](env.md) - For configuring environment variables
|
||||
- [Installation Guide](installation.md) - For deployment instructions
|
||||
|
||||
+18
-25
@@ -14,20 +14,18 @@ Memory is the first bottleneck; CPU matters only when Celery or the Next.js buil
|
||||
|
||||
## 2. Development Environment Memory Requirements
|
||||
|
||||
| Service | Typical use | Rationale / source |
|
||||
|------------------|-------------------------------|-----------------------------------------------------------------------------------------|
|
||||
| PostgreSQL | **1 – 2 GB** | `shared_buffers` starting point ≈ 25% RAM ([postgresql.org][1]) |
|
||||
| Keycloak | **≈ 1.3 GB** | 70% of limit for heap + ~300 MB non-heap ([keycloak.org][2]) |
|
||||
| Redis | **≤ 256 MB** | Empty instance ≈ 3 MB; budget 256 MB to allow small datasets ([stackoverflow.com][3]) |
|
||||
| MinIO | **2 GB (dev) / 32 GB (prod)** | Pre-allocates 1–2 GiB; docs recommend 32 GB per host for ≤ 100 Ti storage ([min.io][4]) |
|
||||
| Django API | **0.8 – 1.5 GB** | Empirical in-house metrics |
|
||||
| Next.js frontend | **0.5 – 1 GB** | Dev build chain |
|
||||
| Nginx | **< 100 MB** | Static reverse-proxy footprint |
|
||||
| Service | Typical use | Rationale / source |
|
||||
|-----------------------|-------------------------------|-----------------------------------------------------------------------------------------|
|
||||
| PostgreSQL | **1 – 2 GB** | `shared_buffers` starting point ≈ 25% RAM ([postgresql.org][1]) |
|
||||
| Keycloak | **≈ 1.3 GB** | 70% of limit for heap + ~300 MB non-heap ([keycloak.org][2]) |
|
||||
| Redis | **≤ 256 MB** | Empty instance ≈ 3 MB; budget 256 MB to allow small datasets ([stackoverflow.com][3]) |
|
||||
| Django API (+ Celery) | **0.8 – 1.5 GB** | Empirical in-house metrics |
|
||||
| Next.js frontend | **0.5 – 1 GB** | Dev build chain |
|
||||
| Nginx | **< 100 MB** | Static reverse-proxy footprint |
|
||||
|
||||
[1]: https://www.postgresql.org/docs/9.1/runtime-config-resource.html "PostgreSQL: Documentation: 9.1: Resource Consumption"
|
||||
[2]: https://www.keycloak.org/high-availability/concepts-memory-and-cpu-sizing "Concepts for sizing CPU and memory resources - Keycloak"
|
||||
[3]: https://stackoverflow.com/questions/45233052/memory-footprint-for-redis-empty-instance "Memory footprint for Redis empty instance - Stack Overflow"
|
||||
[4]: https://min.io/docs/minio/kubernetes/upstream/operations/checklists/hardware.html "Hardware Checklist — MinIO Object Storage for Kubernetes"
|
||||
|
||||
> **Rule of thumb:** add 2 GB for OS/overhead, then sum only the rows you actually run.
|
||||
|
||||
@@ -40,7 +38,6 @@ Production deployments differ significantly from development environments. The t
|
||||
| PostgreSQL | **2 – 8 GB** | Higher `shared_buffers` and connection pooling for concurrent users |
|
||||
| OIDC Provider (optional) | **Variable** | Any OIDC-compatible provider (Keycloak, Auth0, Azure AD, etc.) - external or self-hosted |
|
||||
| Redis | **256 MB – 2 GB** | Session storage and caching; scales with active user sessions |
|
||||
| Object Storage (optional) | **External or self-hosted** | Can use AWS S3, Azure Blob, Google Cloud Storage, or self-hosted MinIO |
|
||||
| Django API (+ Celery) | **1 – 3 GB** | Production workloads with background tasks and higher concurrency |
|
||||
| Static Files (Nginx) | **< 200 MB** | Serves Next.js build output and static assets; no development overhead |
|
||||
| Nginx (Load Balancer) | **< 200 MB** | Reverse proxy, SSL termination, static file serving |
|
||||
@@ -49,7 +46,6 @@ Production deployments differ significantly from development environments. The t
|
||||
|
||||
- **Frontend**: Uses pre-built Next.js static assets served by Nginx (no Node.js runtime needed)
|
||||
- **Authentication**: Any OIDC-compatible provider can be used instead of self-hosted Keycloak
|
||||
- **Object Storage**: External services (S3, Azure Blob) or self-hosted solutions (MinIO) are both viable
|
||||
- **Database**: Consider PostgreSQL clustering or managed database services for high availability
|
||||
- **Scaling**: Horizontal scaling is recommended for Django API service
|
||||
|
||||
@@ -58,7 +54,7 @@ Production deployments differ significantly from development environments. The t
|
||||
| Service | Memory | Notes |
|
||||
|----------------------------------|------------|----------------------------------------|
|
||||
| PostgreSQL | **2 GB** | Core database |
|
||||
| Django API | **1.5 GB** | Backend services |
|
||||
| Django API (+ Celery) | **1.5 GB** | Backend services |
|
||||
| Nginx | **100 MB** | Static files + reverse proxy |
|
||||
| Redis | **256 MB** | Session storage |
|
||||
| **Total (without auth/storage)** | **≈ 4 GB** | External OIDC + object storage assumed |
|
||||
@@ -81,16 +77,15 @@ Production deployments differ significantly from development environments. The t
|
||||
|
||||
## 5. Ports (dev defaults)
|
||||
|
||||
| Port | Service |
|
||||
|-----------|----------------------------|
|
||||
| 3000 | Next.js |
|
||||
| 8071 | Django |
|
||||
| 8080 | Keycloak |
|
||||
| 8083 | Nginx proxy |
|
||||
| 9000/9001 | MinIO |
|
||||
| 15432 | PostgreSQL (main) |
|
||||
| 5433 | PostgreSQL (Keycloak) |
|
||||
| 1081 | Maildev (currently unused) |
|
||||
| Port | Service |
|
||||
|-----------|-----------------------|
|
||||
| 3000 | Next.js |
|
||||
| 8071 | Django |
|
||||
| 8080 | Keycloak |
|
||||
| 8083 | Nginx proxy |
|
||||
| 15432 | PostgreSQL (main) |
|
||||
| 5433 | PostgreSQL (Keycloak) |
|
||||
| 1081 | Maildev |
|
||||
|
||||
## 6. Sizing Guidelines
|
||||
|
||||
@@ -101,5 +96,3 @@ Production deployments differ significantly from development environments. The t
|
||||
**CPU** – budget one vCPU per busy container until Celery or Next.js builds saturate.
|
||||
|
||||
**Disk** – SSD; add 10 GB extra for the Docker layer cache.
|
||||
|
||||
**MinIO** – for demos, mount a local folder instead of running MinIO to save 2 GB+ of RAM.
|
||||
|
||||
+4
-4
@@ -4,7 +4,7 @@
|
||||
|
||||
To use this feature, simply set the `FRONTEND_CSS_URL` environment variable to the URL of your custom CSS file. For example:
|
||||
|
||||
```ini
|
||||
```javascript
|
||||
FRONTEND_CSS_URL=http://anything/custom-style.css
|
||||
```
|
||||
|
||||
@@ -38,7 +38,7 @@ The footer is configurable from the theme customization file.
|
||||
|
||||
### Settings 🔧
|
||||
|
||||
```ini
|
||||
```shellscript
|
||||
THEME_CUSTOMIZATION_FILE_PATH=<path>
|
||||
```
|
||||
|
||||
@@ -55,10 +55,10 @@ The translations can be partially overridden from the theme customization file.
|
||||
|
||||
### Settings 🔧
|
||||
|
||||
```ini
|
||||
```shellscript
|
||||
THEME_CUSTOMIZATION_FILE_PATH=<path>
|
||||
```
|
||||
|
||||
### Example of JSON
|
||||
|
||||
The json must follow some rules: https://github.com/suitenumerique/conversations/blob/main/src/helm/env.d/dev/configuration/theme/demo.json
|
||||
The json must follow some rules: https://github.com/suitenumerique/conversations/blob/main/src/helm/env.d/dev/configuration/theme/demo.json
|
||||
-238
@@ -1,238 +0,0 @@
|
||||
# Tools for the Conversation Agent
|
||||
|
||||
The conversation agent can be extended with various tools that provide additional capabilities such as web search,
|
||||
weather information, and more. We currently only have web search tools, but more tools can be added as needed.
|
||||
This document explains how to configure and use these tools.
|
||||
|
||||
## Overview
|
||||
|
||||
Tools are functions that the LLM can call during a conversation to access external data or perform specific actions.
|
||||
The agent decides when to use these tools based on the user's query and the conversation context.
|
||||
|
||||
## Configuring Tools for a Model
|
||||
|
||||
Tools are configured at the model level in the LLM configuration file.
|
||||
Each model can have its own set of available tools.
|
||||
|
||||
### Configuration File Location
|
||||
|
||||
Read the [LLM Configuration](llm-configuration.md) document to find out where the configuration file is located
|
||||
and how to use it.
|
||||
|
||||
### Example Configuration
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "default-model",
|
||||
"model_name": "gpt-4",
|
||||
"human_readable_name": "GPT-4 with Tools",
|
||||
"provider_name": "default-provider",
|
||||
"is_active": true,
|
||||
"system_prompt": "You are a helpful assistant.",
|
||||
"tools": [
|
||||
"web_search_brave",
|
||||
"get_current_weather"
|
||||
]
|
||||
}
|
||||
],
|
||||
"providers": [
|
||||
{
|
||||
"hrid": "default-provider",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"api_key": "settings.AI_API_KEY",
|
||||
"kind": "openai"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
The `tools` field accepts either:
|
||||
- A list of tool names: `["tool_name_1", "tool_name_2"]`
|
||||
- A reference to a settings variable: `"settings.AI_AGENT_TOOLS"`
|
||||
|
||||
## Available Tools
|
||||
|
||||
To make a tool available to be in a model's configuration, it must be registered in the tool registry located at
|
||||
`src/backend/chat/tools/__init__.py`.
|
||||
|
||||
This is not dynamic - any changes to the tool registry require a code deployment...
|
||||
We want to add dynamic loading in the future.
|
||||
|
||||
| Tool Name | Description | Documentation |
|
||||
|------------------------------------------|---------------------------------------------------------------|-----------------------------------------------------------------------------|
|
||||
| `get_current_weather` | Fake weather tool for testing purposes | [Details](tools/get_current_weather.md) |
|
||||
| `web_search_tavily` | Web search using Tavily API | [Details](tools/web_search_tavily.md) |
|
||||
| `web_search_brave` | Web search using Brave Search API with optional summarization | [Details](tools/web_search_brave.md) |
|
||||
| `web_search_brave_with_document_backend` | Web search using Brave with RAG-based document processing | [Details](tools/web_search_brave.md#web_search_brave_with_document_backend) |
|
||||
| `web_search_albert_rag` | ⚠️ **Deprecated** - Web search using Albert API with RAG | [Details](tools/web_search_brave.md#deprecated-web_search_albert_rag) |
|
||||
|
||||
## Adding a New Tool
|
||||
|
||||
To add a new tool to the system, follow these steps:
|
||||
|
||||
### 1. Create the Tool Function
|
||||
|
||||
Create a new Python file in `src/backend/chat/tools/` with your tool function. The function should:
|
||||
|
||||
- Have clear type annotations
|
||||
- Include a comprehensive docstring (the LLM uses this to understand when to use the tool)
|
||||
- Accept `RunContext` as the first parameter if it needs access to conversation context
|
||||
- Return appropriate data types
|
||||
|
||||
Example:
|
||||
```python
|
||||
"""My custom tool for the chat agent."""
|
||||
|
||||
from pydantic_ai import RunContext
|
||||
|
||||
def my_custom_tool(ctx: RunContext, param1: str, param2: int) -> dict:
|
||||
"""
|
||||
Brief description of what the tool does.
|
||||
|
||||
The LLM uses this description to decide when to call this tool.
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
param1 (str): Description of parameter 1.
|
||||
param2 (int): Description of parameter 2.
|
||||
|
||||
Returns:
|
||||
dict: Description of the return value.
|
||||
"""
|
||||
# Your implementation here
|
||||
return {"result": "example"}
|
||||
```
|
||||
|
||||
### 2. Register the Tool
|
||||
|
||||
Add your tool to the registry in `src/backend/chat/tools/__init__.py`:
|
||||
|
||||
```python
|
||||
from .my_custom_tool import my_custom_tool
|
||||
|
||||
def get_pydantic_tools_by_name(name: str) -> Tool:
|
||||
"""Get a tool by its name."""
|
||||
tool_dict = {
|
||||
"get_current_weather": Tool(get_current_weather, takes_ctx=False),
|
||||
"web_search_brave": Tool(
|
||||
web_search_brave, takes_ctx=False, prepare=only_if_web_search_enabled
|
||||
),
|
||||
# Add your tool here
|
||||
"my_custom_tool": Tool(
|
||||
my_custom_tool,
|
||||
takes_ctx=True, # Set to True if your tool needs RunContext
|
||||
# prepare=only_if_web_search_enabled # Optional: add conditions
|
||||
),
|
||||
}
|
||||
return tool_dict[name]
|
||||
```
|
||||
|
||||
### 3. Update Imports
|
||||
|
||||
Don't forget to import your tool function at the top of `__init__.py`:
|
||||
|
||||
```python
|
||||
from .my_custom_tool import my_custom_tool
|
||||
```
|
||||
|
||||
### 4. Add to Model Configuration
|
||||
|
||||
Add your tool name to the `tools` list in your LLM configuration file or
|
||||
to the `AI_AGENT_TOOLS` environment variable for local/test purpose.
|
||||
|
||||
## Tool Preparation: Conditional Tool Availability
|
||||
|
||||
Some tools should only be available under certain conditions. The `prepare` parameter in the `Tool` constructor
|
||||
allows you to specify a function that determines whether a tool should be included.
|
||||
|
||||
### The `only_if_web_search_enabled` Prepare Function
|
||||
|
||||
This is a built-in prepare function that checks if web search feature is enabled in the conversation context:
|
||||
|
||||
```python
|
||||
async def only_if_web_search_enabled(ctx, tool_def: ToolDefinition) -> ToolDefinition | None:
|
||||
"""Prepare function to include a tool only if web search is enabled in the context."""
|
||||
return tool_def if ctx.deps.web_search_enabled else None
|
||||
```
|
||||
|
||||
### Usage
|
||||
|
||||
All web search tools use this prepare function:
|
||||
|
||||
```python
|
||||
"web_search_brave": Tool(
|
||||
web_search_brave,
|
||||
takes_ctx=False,
|
||||
prepare=only_if_web_search_enabled
|
||||
),
|
||||
```
|
||||
|
||||
This ensures that web search tools are only available when the user or conversation settings have enabled web search functionality.
|
||||
|
||||
### Creating Custom Prepare Functions
|
||||
|
||||
You can create your own prepare functions for custom conditions:
|
||||
|
||||
```python
|
||||
async def only_if_feature_enabled(ctx, tool_def: ToolDefinition) -> ToolDefinition | None:
|
||||
"""Include tool only if a specific feature is enabled."""
|
||||
return tool_def if ctx.deps.feature_enabled else None
|
||||
```
|
||||
|
||||
## Web Search Enable/Disable
|
||||
|
||||
Web search tools can be toggled on or off based on conversation settings. When web search is disabled:
|
||||
- Web search tools are not included in the agent's available tools
|
||||
- The LLM cannot make web search calls even if it tries
|
||||
- This is enforced by the `only_if_web_search_enabled` prepare function
|
||||
|
||||
The `web_search_enabled` flag is typically set:
|
||||
- Per conversation in the conversation settings
|
||||
- Per user preference
|
||||
- Through admin configuration
|
||||
|
||||
## Best Practices
|
||||
|
||||
1. **Keep tools focused** - Each tool should do one thing well
|
||||
2. **Clear documentation** - The LLM relies on docstrings to understand when to use tools
|
||||
3. **Error handling** - Tools should handle errors gracefully and return meaningful messages
|
||||
4. **Performance** - Be mindful of API rate limits and timeout values
|
||||
5. **Security** - Never log sensitive data (API keys, user data, etc.)
|
||||
6. **Caching** - Use Django's cache framework for expensive operations when appropriate
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Tool Not Being Called
|
||||
|
||||
If the LLM isn't calling your tool:
|
||||
- Check that the tool is registered in `get_pydantic_tools_by_name`
|
||||
- Verify the tool is in the model's `tools` configuration
|
||||
- Review the tool's docstring - make it clearer when the tool should be used
|
||||
- Check if any `prepare` function is preventing the tool from being included
|
||||
|
||||
### Tool Errors
|
||||
|
||||
If a tool is throwing errors:
|
||||
- Check the logs for detailed error messages
|
||||
- Verify all required environment variables are set
|
||||
- Ensure the tool's dependencies are installed
|
||||
- Test the tool function independently
|
||||
|
||||
We recommend wrapping external API calls in try/except blocks to handle potential issues gracefully and use
|
||||
the Pydantic AI `ModelRetry` exception to let the LLM manage the errors.
|
||||
|
||||
### Tool Response Issues
|
||||
|
||||
If the LLM isn't using the tool response correctly:
|
||||
- Ensure the return type is clear and well-structured
|
||||
- Consider returning a `ToolReturn` object with metadata
|
||||
- Check if the response format matches what the LLM expects
|
||||
|
||||
## See Also
|
||||
|
||||
- [Web Search Configuration](llm-configuration.md)
|
||||
- [Architecture](architecture.md)
|
||||
- [Environment Variables](env.md)
|
||||
|
||||
@@ -1,113 +0,0 @@
|
||||
# get_current_weather Tool
|
||||
|
||||
## Overview
|
||||
|
||||
The `get_current_weather` tool is a **fake weather tool** designed for testing and demonstration purposes. It does not connect to any real weather API and always returns hardcoded weather data.
|
||||
|
||||
## Purpose
|
||||
|
||||
This tool is useful for:
|
||||
- **Testing** the tool calling functionality of LLMs
|
||||
- **Demonstrating** how tools work without requiring API keys
|
||||
- **Development** and debugging of the agent system
|
||||
- **Example implementation** for creating new tools
|
||||
|
||||
⚠️ **Warning**: This tool should **not** be used in production environments. It always returns fake data regardless of the location or conditions.
|
||||
|
||||
## Configuration
|
||||
|
||||
### Add to Model
|
||||
|
||||
To enable this tool for a model, add it to the `tools` list in your LLM configuration:
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "my-model",
|
||||
"tools": [
|
||||
"get_current_weather"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Or via environment variable when using local environment settings:
|
||||
```ini
|
||||
AI_AGENT_TOOLS=get_current_weather
|
||||
```
|
||||
|
||||
### No Additional Settings Required
|
||||
|
||||
This tool does not require any API keys, environment variables, or additional configuration.
|
||||
|
||||
## Function Signature
|
||||
|
||||
```python
|
||||
def get_current_weather(location: str, unit: str) -> dict:
|
||||
"""
|
||||
Get the current weather in a given location.
|
||||
|
||||
Args:
|
||||
location (str): The city and state, e.g. San Francisco, CA.
|
||||
unit (str): The unit of temperature, either 'celsius' or 'fahrenheit'.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the location, temperature, and unit.
|
||||
"""
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|------------|------|----------|-----------------------------------------------------------------|
|
||||
| `location` | str | Yes | The city and state (e.g., "San Francisco, CA", "Paris, France") |
|
||||
| `unit` | str | Yes | Temperature unit: either "celsius" or "fahrenheit" |
|
||||
|
||||
## Return Value
|
||||
|
||||
Returns a dictionary with the following structure:
|
||||
|
||||
```python
|
||||
{
|
||||
"location": str, # The location that was queried
|
||||
"temperature": int, # Always 22°C or 72°F
|
||||
"unit": str # The unit that was requested
|
||||
}
|
||||
```
|
||||
|
||||
## How the LLM Uses It
|
||||
|
||||
When a user asks about weather, the LLM will:
|
||||
|
||||
1. **Recognize** the weather-related query
|
||||
2. **Extract** the location from the user's message
|
||||
3. **Determine** the appropriate unit (often from context or user preference)
|
||||
4. **Call** the `get_current_weather` tool
|
||||
5. **Receive** the fake weather data
|
||||
6. **Format** a response to the user
|
||||
|
||||
### Example Conversation
|
||||
|
||||
**User**: "What's the weather like in London?"
|
||||
|
||||
**LLM** (internal): *Calls `get_current_weather("London, UK", "celsius")`*
|
||||
|
||||
**Tool Response**:
|
||||
```json
|
||||
{
|
||||
"location": "London, UK",
|
||||
"temperature": 22,
|
||||
"unit": "celsius"
|
||||
}
|
||||
```
|
||||
|
||||
**LLM** (to user): "The current weather in London, UK is 22°C."
|
||||
|
||||
## See Also
|
||||
|
||||
- [Tools Overview](../tools.md)
|
||||
- [Adding a New Tool](../tools.md#adding-a-new-tool)
|
||||
- [Testing Tools](../tools.md#testing-your-tools)
|
||||
|
||||
@@ -1,670 +0,0 @@
|
||||
# Brave Web Search Tools
|
||||
|
||||
## Overview
|
||||
|
||||
The Brave web search tools enable the conversation agent to search the web using the [Brave Search API](https://brave.com/search/api/).
|
||||
Brave Search is a privacy-focused search engine that provides comprehensive web search results.
|
||||
|
||||
This documentation covers three related tools:
|
||||
1. **`web_search_brave`** - Standard web search with optional summarization
|
||||
2. **`web_search_brave_with_document_backend`** - Web search with RAG-based document processing
|
||||
3. **`web_search_albert_rag`** - ⚠️ **Deprecated** - Use `web_search_brave_with_document_backend` instead
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Common Configuration](#common-configuration)
|
||||
- [web_search_brave](#web_search_brave)
|
||||
- [web_search_brave_with_document_backend](#web_search_brave_with_document_backend)
|
||||
- [Deprecated: web_search_albert_rag](#deprecated-web_search_albert_rag)
|
||||
- [Comparison](#comparison)
|
||||
- [Best Practices](#best-practices)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
|
||||
---
|
||||
|
||||
## Common Configuration
|
||||
|
||||
### Prerequisites
|
||||
|
||||
1. **Brave Search API Key**: Sign up at [Brave Search API](https://brave.com/search/api/) to get an API key
|
||||
2. **Environment Variables**: Configure the required settings
|
||||
|
||||
### Common Environment Variables
|
||||
|
||||
All Brave tools share these common settings:
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|---------------------|----------|---------|----------------------------------------------------|
|
||||
| `BRAVE_API_KEY` | **Yes** | None | Your Brave Search API key |
|
||||
| `BRAVE_API_TIMEOUT` | No | 5 | API request timeout in seconds |
|
||||
| `BRAVE_MAX_RESULTS` | No | 8 | Maximum number of search results |
|
||||
| `BRAVE_CACHE_TTL` | No | 1800 | Cache time-to-live in seconds (30 minutes) |
|
||||
|
||||
### Search Parameters
|
||||
|
||||
Check on the Brave API documentation for more details on these parameters:
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|-------------------------------|----------|------------|---------------------------------------------------|
|
||||
| `BRAVE_SEARCH_COUNTRY` | No | None | Country code for search (e.g., "US", "FR") |
|
||||
| `BRAVE_SEARCH_LANG` | No | None | Language code (e.g., "en", "fr") |
|
||||
| `BRAVE_SEARCH_SAFE_SEARCH` | No | "moderate" | Safe search level: "off", "moderate", or "strict" |
|
||||
| `BRAVE_SEARCH_SPELLCHECK` | No | True | Enable spell checking |
|
||||
| `BRAVE_SEARCH_EXTRA_SNIPPETS` | No | True | Fetch extra snippets from pages |
|
||||
|
||||
|
||||
Note: even if `BRAVE_SEARCH_EXTRA_SNIPPETS` is enabled, the API may not include them if you don't have a plan for this.
|
||||
This is why, in `web_search_brave`, we also fetch the page content ourselves when needed.
|
||||
|
||||
### Configuration Example
|
||||
|
||||
```bash
|
||||
# .env file
|
||||
BRAVE_API_KEY=BSA-your-api-key-here
|
||||
BRAVE_MAX_RESULTS=8
|
||||
BRAVE_MAX_WORKERS=4
|
||||
BRAVE_SEARCH_COUNTRY=US
|
||||
BRAVE_SEARCH_LANG=en
|
||||
BRAVE_SEARCH_SAFE_SEARCH=moderate
|
||||
```
|
||||
|
||||
### Django Settings
|
||||
|
||||
All Brave settings are defined in `src/backend/conversations/brave_settings.py`:
|
||||
|
||||
```python
|
||||
class BraveSettings:
|
||||
"""Brave settings for web_search_brave tool."""
|
||||
|
||||
BRAVE_API_KEY = values.Value(
|
||||
default=None,
|
||||
environ_name="BRAVE_API_KEY",
|
||||
environ_prefix=None,
|
||||
)
|
||||
# ... more settings
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## web_search_brave
|
||||
|
||||
### Overview
|
||||
|
||||
Standard Brave web search tool with optional LLM-based summarization of page content.
|
||||
|
||||
### Purpose
|
||||
|
||||
- Search the web for up-to-date information
|
||||
- Extract content from web pages
|
||||
- Optionally summarize content using an LLM
|
||||
- Provide structured results with snippets
|
||||
|
||||
### Additional Configuration
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|-------------------------------|----------|---------|-------------------------------------------------|
|
||||
| `BRAVE_SUMMARIZATION_ENABLED` | No | False | Enable LLM-based summarization of fetched pages |
|
||||
|
||||
### Function Signature
|
||||
|
||||
```python
|
||||
def web_search_brave(query: str) -> ToolReturn:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
query (str): The query to search for.
|
||||
|
||||
Returns:
|
||||
ToolReturn: Formatted search results with metadata
|
||||
"""
|
||||
```
|
||||
|
||||
### Return Value
|
||||
|
||||
Returns a `ToolReturn` object with:
|
||||
|
||||
```python
|
||||
ToolReturn(
|
||||
return_value={
|
||||
"0": {
|
||||
"url": "https://example.com/page1",
|
||||
"title": "Example Page Title",
|
||||
"snippets": ["Extracted or summarized content..."]
|
||||
},
|
||||
"1": {
|
||||
"url": "https://example.com/page2",
|
||||
"title": "Another Page",
|
||||
"snippets": ["More content..."]
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"sources": {
|
||||
"https://example.com/page1",
|
||||
"https://example.com/page2"
|
||||
}
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
### How It Works
|
||||
|
||||
1. **Query API**: Sends search query to Brave Search API
|
||||
2. **Receive Results**: Gets list of matching web pages
|
||||
3. **Fetch Content**: For results without extra_snippets:
|
||||
- Fetches the HTML content using `trafilatura`
|
||||
- Extracts the main text content
|
||||
- Caches the extracted content
|
||||
4. **Summarize (Optional)**: If `BRAVE_SUMMARIZATION_ENABLED=True`:
|
||||
- Sends extracted content to summarization agent
|
||||
- Receives concise summary focused on the query
|
||||
5. **Format Results**: Returns structured data with URLs, titles, and snippets
|
||||
|
||||
### Workflow Diagram
|
||||
|
||||
```
|
||||
User Query
|
||||
↓
|
||||
Brave Search API
|
||||
↓
|
||||
Search Results (URLs, titles, descriptions)
|
||||
↓
|
||||
[For each result without snippets]
|
||||
↓
|
||||
Fetch HTML (trafilatura) → Extract Text → Cache
|
||||
↓
|
||||
[If BRAVE_SUMMARIZATION_ENABLED]
|
||||
↓
|
||||
Summarization Agent (LLM)
|
||||
↓
|
||||
Summary Text
|
||||
↓
|
||||
Format & Return
|
||||
```
|
||||
|
||||
### Caching
|
||||
|
||||
Extracted content is cached to avoid repeated fetches:
|
||||
|
||||
```python
|
||||
cache_key = f"web_search_brave:extract:{url}"
|
||||
cache.set(cache_key, document, settings.BRAVE_CACHE_TTL)
|
||||
```
|
||||
|
||||
**Cache Duration**: Controlled by `BRAVE_CACHE_TTL` (default: 30 minutes)
|
||||
|
||||
### Summarization
|
||||
|
||||
When enabled, the tool uses the `SummarizationAgent` to condense page content:
|
||||
|
||||
```python
|
||||
prompt = f"""
|
||||
Based on the following request, summarize the following text in a concise manner,
|
||||
focusing on the key points regarding the user request.
|
||||
The result should be up to 30 lines long.
|
||||
|
||||
<user request>
|
||||
{query}
|
||||
</user request>
|
||||
|
||||
<text to summarize>
|
||||
{text}
|
||||
</text to summarize>
|
||||
"""
|
||||
```
|
||||
|
||||
**Note**: Summarization is costly (additional LLM calls).
|
||||
Use only when necessary, we prefer the document vector search from `web_search_brave_with_document_backend`.
|
||||
|
||||
### Add to Model
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "my-model",
|
||||
"tools": [
|
||||
"web_search_brave"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Example Usage
|
||||
|
||||
**User**: "What are the new features in Django 5.0?"
|
||||
|
||||
**Tool Call**: `web_search_brave("Django 5.0 new features")`
|
||||
|
||||
**Tool Response**:
|
||||
```python
|
||||
{
|
||||
"0": {
|
||||
"url": "https://docs.djangoproject.com/en/5.0/releases/5.0/",
|
||||
"title": "Django 5.0 release notes",
|
||||
"snippets": ["Django 5.0 introduces several new features including..."]
|
||||
},
|
||||
# ... more results
|
||||
}
|
||||
```
|
||||
|
||||
### Registration
|
||||
|
||||
```python
|
||||
"web_search_brave": Tool(
|
||||
web_search_brave,
|
||||
takes_ctx=False,
|
||||
prepare=only_if_web_search_enabled
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## web_search_brave_with_document_backend
|
||||
|
||||
### Overview
|
||||
|
||||
Advanced Brave web search tool that uses RAG (Retrieval-Augmented Generation)
|
||||
with a document backend for intelligent content processing and retrieval.
|
||||
|
||||
### Purpose
|
||||
|
||||
- Search the web and process results through a RAG system
|
||||
- Store fetched documents in a temporary vector database
|
||||
- Perform semantic search across fetched content
|
||||
- Return the most relevant chunks based on the query
|
||||
|
||||
### Additional Configuration
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|-------------------------------------|----------|------------------|----------------------------------------------|
|
||||
| `BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER` | No | 10 | Number of chunks to retrieve from RAG search |
|
||||
| `RAG_DOCUMENT_SEARCH_BACKEND` | No | AlbertRagBackend | Document backend for RAG processing |
|
||||
|
||||
### Function Signature
|
||||
|
||||
```python
|
||||
def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The query to search for.
|
||||
|
||||
Returns:
|
||||
ToolReturn: Formatted search results with RAG-enhanced snippets
|
||||
"""
|
||||
```
|
||||
|
||||
### How It Works
|
||||
|
||||
1. **Query API**: Sends search query to Brave Search API
|
||||
2. **Receive Results**: Gets list of matching web pages
|
||||
3. **Create Temporary Collection**: Creates a temporary vector database collection
|
||||
4. **Fetch & Store**: For each result:
|
||||
- Fetches the HTML content
|
||||
- Extracts the main text
|
||||
- Stores in the temporary document backend
|
||||
5. **RAG Search**: Performs semantic search across stored documents
|
||||
6. **Map Results**: Maps RAG chunks back to original search results
|
||||
7. **Format & Return**: Returns structured data with enhanced snippets
|
||||
8. **Cleanup**: Temporary collection is automatically deleted
|
||||
|
||||
### Workflow Diagram
|
||||
|
||||
```
|
||||
User Query
|
||||
↓
|
||||
Brave Search API
|
||||
↓
|
||||
Search Results (URLs)
|
||||
↓
|
||||
Create Temporary Vector Collection
|
||||
↓
|
||||
[For each URL]
|
||||
↓
|
||||
Fetch HTML → Extract Text → Store in Vector DB
|
||||
↓
|
||||
RAG Semantic Search
|
||||
↓
|
||||
Retrieve Most Relevant Chunks
|
||||
↓
|
||||
Map Chunks to Original URLs
|
||||
↓
|
||||
Format & Return
|
||||
↓
|
||||
Delete Temporary Collection
|
||||
```
|
||||
|
||||
### Temporary Collection
|
||||
|
||||
The tool creates a temporary collection with a unique ID:
|
||||
|
||||
```python
|
||||
with document_store_backend.temporary_collection(f"tmp-{uuid.uuid4()}") as document_store:
|
||||
# Fetch and store documents
|
||||
# Perform search
|
||||
# Collection is automatically deleted on exit
|
||||
```
|
||||
|
||||
### RAG Search
|
||||
|
||||
The RAG backend performs semantic search to find the most relevant content:
|
||||
|
||||
```python
|
||||
rag_results = document_store.search(
|
||||
query,
|
||||
results_count=settings.BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER,
|
||||
)
|
||||
```
|
||||
|
||||
Returns chunks ranked by relevance to the query, not just keyword matching.
|
||||
|
||||
### Token Usage Tracking
|
||||
|
||||
The tool tracks LLM tokens used during RAG processing:
|
||||
|
||||
```python
|
||||
ctx.usage += RunUsage(
|
||||
input_tokens=rag_results.usage.prompt_tokens,
|
||||
output_tokens=rag_results.usage.completion_tokens,
|
||||
)
|
||||
```
|
||||
|
||||
### Document Backend
|
||||
|
||||
The default backend is `AlbertRagBackend`, but you can configure a different one:
|
||||
|
||||
```bash
|
||||
RAG_DOCUMENT_SEARCH_BACKEND=chat.agent_rag.document_rag_backends.custom_backend.CustomBackend
|
||||
```
|
||||
|
||||
### Add to Model
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "my-model",
|
||||
"tools": [
|
||||
"web_search_brave_with_document_backend"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Example Usage
|
||||
|
||||
**User**: "Explain the concept of async views in Django"
|
||||
|
||||
**Tool Call**: `web_search_brave_with_document_backend(ctx, "Django async views explained")`
|
||||
|
||||
**Tool Response**:
|
||||
```python
|
||||
{
|
||||
"0": {
|
||||
"url": "https://docs.djangoproject.com/en/stable/topics/async/",
|
||||
"title": "Asynchronous support",
|
||||
"snippets": [
|
||||
"Django has support for writing asynchronous views...",
|
||||
"Async views are declared using Python's async def syntax..."
|
||||
]
|
||||
},
|
||||
# ... more results with relevant chunks
|
||||
}
|
||||
```
|
||||
|
||||
### Registration
|
||||
|
||||
```python
|
||||
"web_search_brave_with_document_backend": Tool(
|
||||
web_search_brave_with_document_backend,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
)
|
||||
```
|
||||
|
||||
### Advantages Over Standard web_search_brave
|
||||
|
||||
| Feature | web_search_brave | web_search_brave_with_document_backend |
|
||||
|-------------------|--------------------------------|----------------------------------------|
|
||||
| Content Retrieval | Full page or summary | Semantic chunks |
|
||||
| Relevance | Keyword-based | Semantic similarity |
|
||||
| Token Efficiency | May include irrelevant content | Only relevant chunks |
|
||||
| Processing | Simpler, faster | More intelligent, slower |
|
||||
| Cost | Lower | Higher (RAG processing) |
|
||||
| Best For | General search | Deep research, technical queries |
|
||||
|
||||
---
|
||||
|
||||
## Deprecated: web_search_albert_rag
|
||||
|
||||
### ⚠️ Deprecation Notice
|
||||
|
||||
The `web_search_albert_rag` tool is **deprecated** and should not be used in new implementations.
|
||||
|
||||
**Replacement**: Use `web_search_brave_with_document_backend` instead, which provides:
|
||||
- Better performance
|
||||
- More control over the RAG backend
|
||||
- Temporary collections (no cleanup issues)
|
||||
- Token usage tracking
|
||||
- Parallel processing support
|
||||
|
||||
### Why Deprecated?
|
||||
|
||||
- Limited to Albert API only
|
||||
- No control over document backend
|
||||
- Less flexible than the new approach
|
||||
- Maintenance burden
|
||||
|
||||
### Timeline
|
||||
|
||||
- **Current**: Still functional but not recommended
|
||||
- **Future**: Will be removed in a future version
|
||||
|
||||
---
|
||||
|
||||
## Comparison
|
||||
|
||||
### When to Use Which Tool?
|
||||
|
||||
#### Use `web_search_brave`
|
||||
|
||||
✅ **Best for**:
|
||||
- General web search queries
|
||||
- Quick information retrieval
|
||||
- When speed is important
|
||||
- Lower cost requirements
|
||||
- Simple fact-finding
|
||||
|
||||
❌ **Not ideal for**:
|
||||
- Deep research requiring precise context
|
||||
- Technical documentation queries
|
||||
- When semantic relevance is crucial
|
||||
|
||||
#### Use `web_search_brave_with_document_backend`
|
||||
|
||||
✅ **Best for**:
|
||||
- Complex technical queries
|
||||
- Research requiring precise context
|
||||
- When semantic relevance is important
|
||||
- Questions needing deep understanding
|
||||
- Documentation and how-to queries
|
||||
|
||||
❌ **Not ideal for**:
|
||||
- Simple factual queries
|
||||
- When speed is critical
|
||||
- Budget-constrained scenarios
|
||||
- High-volume usage
|
||||
|
||||
---
|
||||
|
||||
## Best Practices
|
||||
|
||||
### Query Formulation
|
||||
|
||||
Help the LLM formulate effective queries:
|
||||
|
||||
```python
|
||||
# Good queries
|
||||
"Python asyncio tutorial 2024"
|
||||
"Django REST framework authentication"
|
||||
"React hooks best practices"
|
||||
|
||||
# Poor queries
|
||||
"tell me about programming" # Too vague
|
||||
"how do I do the thing with the stuff" # Unclear
|
||||
```
|
||||
|
||||
### Performance Optimization
|
||||
|
||||
#### 1. Optimize Cache
|
||||
|
||||
```bash
|
||||
# Longer cache for stable content
|
||||
BRAVE_CACHE_TTL=3600 # 1 hour
|
||||
|
||||
# Shorter cache for dynamic content
|
||||
BRAVE_CACHE_TTL=300 # 5 minutes
|
||||
```
|
||||
|
||||
#### 2. Control Result Count
|
||||
|
||||
```bash
|
||||
# Fewer results = faster responses
|
||||
BRAVE_MAX_RESULTS=5
|
||||
|
||||
# More results = more comprehensive
|
||||
BRAVE_MAX_RESULTS=10
|
||||
```
|
||||
|
||||
### Summarization Best Practices
|
||||
|
||||
Only enable summarization when needed:
|
||||
|
||||
```bash
|
||||
# Enable for long-form content
|
||||
BRAVE_SUMMARIZATION_ENABLED=True
|
||||
|
||||
# Disable for speed
|
||||
BRAVE_SUMMARIZATION_ENABLED=False
|
||||
```
|
||||
|
||||
**Cost consideration**: Summarization makes additional LLM calls for each result,
|
||||
significantly increasing costs (and execution time).
|
||||
|
||||
### RAG Configuration
|
||||
|
||||
For `web_search_brave_with_document_backend`:
|
||||
|
||||
```bash
|
||||
# More chunks = more context, higher cost
|
||||
BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER=10
|
||||
|
||||
# Fewer chunks = faster, less context
|
||||
BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER=5
|
||||
```
|
||||
|
||||
### Search Parameters
|
||||
|
||||
```bash
|
||||
# Localize results
|
||||
BRAVE_SEARCH_COUNTRY=FR
|
||||
BRAVE_SEARCH_LANG=fr
|
||||
|
||||
# Safe search for public deployments
|
||||
BRAVE_SEARCH_SAFE_SEARCH=strict
|
||||
|
||||
# Enable spell check for better results
|
||||
BRAVE_SEARCH_SPELLCHECK=True
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Common Issues
|
||||
|
||||
#### 1. No Results Returned
|
||||
|
||||
**Symptoms**: Empty results or no snippets
|
||||
|
||||
**Causes**:
|
||||
- Query too specific
|
||||
- Content extraction failed
|
||||
- Trafilatura couldn't parse the pages
|
||||
|
||||
**Solutions**:
|
||||
```bash
|
||||
# Enable extra snippets
|
||||
BRAVE_SEARCH_EXTRA_SNIPPETS=True
|
||||
|
||||
# Increase result count
|
||||
BRAVE_MAX_RESULTS=10
|
||||
|
||||
# Check logs for extraction errors
|
||||
```
|
||||
|
||||
#### 2. API Errors
|
||||
|
||||
**Symptoms**: HTTP errors, authentication failures
|
||||
|
||||
**Causes**:
|
||||
- Invalid API key
|
||||
- Rate limit exceeded
|
||||
- API service issues
|
||||
|
||||
**Solutions**:
|
||||
```bash
|
||||
# Verify API key is set
|
||||
echo $BRAVE_API_KEY
|
||||
|
||||
# Check Brave API dashboard for limits
|
||||
# Implement rate limiting in your application
|
||||
```
|
||||
|
||||
#### 3. The tool is not being called
|
||||
**Symptoms**: LLM doesn't use the tool even when appropriate
|
||||
|
||||
**Causes**:
|
||||
- Web search not enabled for the conversation
|
||||
- Tool not in model configuration
|
||||
|
||||
**Solutions**:
|
||||
- Check conversation settings have `web_search_enabled=True`
|
||||
- Verify tool is in the model's `tools` list
|
||||
|
||||
---
|
||||
|
||||
## Security Considerations
|
||||
|
||||
This tool is quite "raw", so be cautious about:
|
||||
- the results returned by the web search
|
||||
- the context size which might be large when not using summarization or RAG if long results are returned
|
||||
- the query content which might include sensitive information
|
||||
- ...
|
||||
|
||||
### Content Validation
|
||||
|
||||
Be aware that fetched content may contain:
|
||||
- Malicious scripts (mitigated by text extraction)
|
||||
- Inappropriate content
|
||||
- Misinformation
|
||||
- Biased information
|
||||
|
||||
The LLM should evaluate sources critically.
|
||||
|
||||
|
||||
---
|
||||
|
||||
## See Also
|
||||
|
||||
- [Tools Overview](../tools.md)
|
||||
- [Tavily Web Search Tool](web_search_tavily.md)
|
||||
- [LLM Configuration](../llm-configuration.md)
|
||||
- [Environment Variables](../env.md)
|
||||
- [Brave Search API Documentation](https://brave.com/search/api/)
|
||||
|
||||
@@ -1,370 +0,0 @@
|
||||
# web_search_tavily Tool
|
||||
|
||||
## Overview
|
||||
|
||||
The `web_search_tavily` tool enables the conversation agent to search the web for up-to-date
|
||||
information using the [Tavily Search API](https://tavily.com/).
|
||||
|
||||
## Purpose
|
||||
|
||||
This tool allows the LLM to:
|
||||
- Access current, real-time information beyond its training data
|
||||
- Answer questions about recent events, news, or developments
|
||||
- Provide factual information with sources
|
||||
- Retrieve specific information from the web
|
||||
|
||||
## Configuration
|
||||
|
||||
### Prerequisites
|
||||
|
||||
1. **Tavily API Key**: Sign up at [Tavily](https://tavily.com/) to get an API key
|
||||
2. **Environment Variables**: Configure the required settings
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Required | Default | Description |
|
||||
|----------------------|----------|---------|--------------------------------------------|
|
||||
| `TAVILY_API_KEY` | **Yes** | None | Your Tavily API key |
|
||||
| `TAVILY_MAX_RESULTS` | No | 5 | Maximum number of search results to return |
|
||||
| `TAVILY_API_TIMEOUT` | No | 10 | API request timeout in seconds |
|
||||
|
||||
### Configuration Example
|
||||
|
||||
```bash
|
||||
# .env file
|
||||
TAVILY_API_KEY=tvly-your-api-key-here
|
||||
TAVILY_MAX_RESULTS=5
|
||||
TAVILY_API_TIMEOUT=10
|
||||
```
|
||||
|
||||
### Add to Model
|
||||
|
||||
To enable this tool for a model, add it to the `tools` list in your LLM configuration:
|
||||
|
||||
```json
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "my-model",
|
||||
"tools": [
|
||||
"web_search_tavily"
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Or via environment variable when using local environment settings:
|
||||
|
||||
```ini
|
||||
AI_AGENT_TOOLS=web_search_tavily
|
||||
```
|
||||
|
||||
## Function Signature
|
||||
|
||||
```python
|
||||
def web_search_tavily(query: str) -> list[dict]:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
query (str): The query to search for.
|
||||
|
||||
Returns:
|
||||
list[dict]: A list of search results, each represented as a dictionary.
|
||||
"""
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------------------|
|
||||
| `query` | str | Yes | The search query string |
|
||||
|
||||
## Return Value
|
||||
|
||||
Returns a list of dictionaries, each containing:
|
||||
|
||||
```python
|
||||
{
|
||||
"link": str, # URL of the result
|
||||
"title": str, # Title of the page
|
||||
"snippet": str # Content snippet from the page
|
||||
}
|
||||
```
|
||||
|
||||
### Example Return Value
|
||||
|
||||
```python
|
||||
[
|
||||
{
|
||||
"link": "https://example.com/article1",
|
||||
"title": "Introduction to Python",
|
||||
"snippet": "Python is a high-level programming language known for its simplicity..."
|
||||
},
|
||||
{
|
||||
"link": "https://example.com/article2",
|
||||
"title": "Python Best Practices",
|
||||
"snippet": "Follow these best practices to write clean and efficient Python code..."
|
||||
}
|
||||
]
|
||||
```
|
||||
|
||||
## How the LLM Uses It
|
||||
|
||||
When a user asks for current information or specific facts:
|
||||
|
||||
1. **LLM recognizes** the need for external information
|
||||
2. **Formulates** an appropriate search query
|
||||
3. **Calls** `web_search_tavily(query="search terms")`
|
||||
4. **Receives** a list of search results
|
||||
5. **Synthesizes** the information into a response
|
||||
6. **Provides** the answer with source references
|
||||
|
||||
### Example Conversation
|
||||
|
||||
**User**: "What are the latest developments in quantum computing?"
|
||||
|
||||
**LLM** (internal): *Calls `web_search_tavily("latest developments quantum computing 2024")`*
|
||||
|
||||
**Tool Response**:
|
||||
```python
|
||||
[
|
||||
{
|
||||
"link": "https://techcrunch.com/quantum-news",
|
||||
"title": "Major Breakthrough in Quantum Computing",
|
||||
"snippet": "Researchers announced a significant breakthrough..."
|
||||
},
|
||||
# ... more results
|
||||
]
|
||||
```
|
||||
|
||||
**LLM** (to user): "Based on recent sources, there have been several developments in quantum computing.
|
||||
Researchers recently announced a breakthrough in error correction. Additionally, new quantum processors
|
||||
with improved qubit stability have been unveiled..."
|
||||
|
||||
## Implementation Details
|
||||
|
||||
### Source Code
|
||||
|
||||
Located at: `src/backend/chat/tools/web_search_tavily.py`
|
||||
|
||||
```python
|
||||
"""Web search tool using Tavily for the chat agent."""
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def web_search_tavily(query: str) -> list[dict]:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
query (str): The query to search for.
|
||||
|
||||
Returns:
|
||||
list[dict]: A list of search results, each represented as a dictionary.
|
||||
"""
|
||||
url = "https://api.tavily.com/search"
|
||||
data = {
|
||||
"query": query,
|
||||
"api_key": settings.TAVILY_API_KEY,
|
||||
"max_results": settings.TAVILY_MAX_RESULTS,
|
||||
}
|
||||
response = requests.post(url, json=data, timeout=settings.TAVILY_API_TIMEOUT)
|
||||
response.raise_for_status()
|
||||
|
||||
json_response = response.json()
|
||||
|
||||
raw_search_results = json_response.get("results", [])
|
||||
|
||||
return [
|
||||
{
|
||||
"link": result["url"],
|
||||
"title": result.get("title", ""),
|
||||
"snippet": result.get("content"),
|
||||
}
|
||||
for result in raw_search_results
|
||||
]
|
||||
```
|
||||
|
||||
### Registration
|
||||
|
||||
The tool is registered in `src/backend/chat/tools/__init__.py`:
|
||||
|
||||
```python
|
||||
"web_search_tavily": Tool(
|
||||
web_search_tavily,
|
||||
takes_ctx=False,
|
||||
prepare=only_if_web_search_enabled
|
||||
)
|
||||
```
|
||||
|
||||
Note that:
|
||||
- `takes_ctx=False` - This tool doesn't need the conversation context
|
||||
- `prepare=only_if_web_search_enabled` - Only available when web search is enabled
|
||||
|
||||
## Django Settings
|
||||
|
||||
The tool uses these Django settings from `settings.py`:
|
||||
|
||||
```python
|
||||
# Tavily API
|
||||
TAVILY_API_KEY = values.Value(
|
||||
None, # Tavily API key is not set by default
|
||||
environ_name="TAVILY_API_KEY",
|
||||
environ_prefix=None,
|
||||
)
|
||||
TAVILY_MAX_RESULTS = values.PositiveIntegerValue(
|
||||
default=5,
|
||||
environ_name="TAVILY_MAX_RESULTS",
|
||||
environ_prefix=None,
|
||||
)
|
||||
TAVILY_API_TIMEOUT = values.PositiveIntegerValue(
|
||||
default=10, # seconds
|
||||
environ_name="TAVILY_API_TIMEOUT",
|
||||
environ_prefix=None,
|
||||
)
|
||||
```
|
||||
|
||||
## Error Handling
|
||||
|
||||
The tool may raise exceptions in the following cases:
|
||||
|
||||
### Missing API Key
|
||||
```python
|
||||
# If TAVILY_API_KEY is not set
|
||||
AttributeError: 'Settings' object has no attribute 'TAVILY_API_KEY'
|
||||
```
|
||||
|
||||
**Solution**: Set the `TAVILY_API_KEY` environment variable
|
||||
|
||||
### API Errors
|
||||
```python
|
||||
# If the API request fails
|
||||
requests.exceptions.HTTPError: 401 Unauthorized
|
||||
```
|
||||
|
||||
**Possible causes**:
|
||||
- Invalid API key
|
||||
- Exceeded rate limits
|
||||
- API service unavailable
|
||||
|
||||
### Timeout Errors
|
||||
```python
|
||||
# If the request takes too long
|
||||
requests.exceptions.Timeout
|
||||
```
|
||||
|
||||
**Solution**: Increase `TAVILY_API_TIMEOUT` or check network connectivity
|
||||
|
||||
## Best Practices
|
||||
|
||||
### Query Formulation
|
||||
|
||||
The LLM should formulate queries that are:
|
||||
- **Specific and focused** - Better results with targeted queries
|
||||
- **Up-to-date** - Include year or "latest" when relevant
|
||||
- **Clear** - Avoid ambiguous terms
|
||||
- **Concise** - Remove unnecessary words
|
||||
|
||||
Good query examples:
|
||||
- ✅ "quantum computing breakthroughs 2024"
|
||||
- ✅ "latest Python 3.12 features"
|
||||
- ✅ "climate change COP29 outcomes"
|
||||
|
||||
Poor query examples:
|
||||
- ❌ "tell me about stuff happening" (too vague)
|
||||
- ❌ "what is the weather like today in Paris on November 5th 2024 at 3pm" (too specific/long)
|
||||
|
||||
### Rate Limiting
|
||||
|
||||
Be aware of Tavily API rate limits:
|
||||
- Free tier: Limited requests per month
|
||||
- Paid tiers: Higher limits
|
||||
|
||||
Monitor your usage and implement caching if needed.
|
||||
|
||||
### Result Count
|
||||
|
||||
The `TAVILY_MAX_RESULTS` setting controls how many results are returned:
|
||||
- **Lower values (3-5)**: Faster responses, less context for LLM
|
||||
- **Higher values (8-10)**: More comprehensive, but slower and more expensive
|
||||
|
||||
Recommended: **5 results** for most use cases
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Tool Not Being Called
|
||||
|
||||
**Symptoms**: LLM doesn't use web search even when appropriate
|
||||
|
||||
**Possible causes**:
|
||||
1. Web search not enabled for the conversation
|
||||
2. Tool not in model configuration
|
||||
3. API key not set
|
||||
|
||||
**Solutions**:
|
||||
1. Check conversation settings have `web_search_enabled=True`
|
||||
2. Verify tool is in the model's `tools` list
|
||||
3. Confirm `TAVILY_API_KEY` is set
|
||||
|
||||
### No Results Returned
|
||||
|
||||
**Symptoms**: Tool returns empty list
|
||||
|
||||
**Possible causes**:
|
||||
1. Query too specific
|
||||
2. No matching results
|
||||
3. API filtering results
|
||||
|
||||
**Solutions**:
|
||||
1. Try broader query terms
|
||||
2. Check Tavily dashboard for query logs
|
||||
3. Review API response in logs
|
||||
|
||||
### Slow Responses
|
||||
|
||||
**Symptoms**: Tool takes a long time to respond
|
||||
|
||||
**Possible causes**:
|
||||
1. Network latency
|
||||
2. Tavily API slow
|
||||
3. Timeout too high
|
||||
|
||||
**Solutions**:
|
||||
1. Check network connectivity
|
||||
2. Monitor Tavily status page
|
||||
3. Adjust `TAVILY_API_TIMEOUT` if needed
|
||||
|
||||
## Security Considerations
|
||||
|
||||
This tool is quite "raw", and was currently only used for test purpose, so be cautious about:
|
||||
- the results returned by the web search
|
||||
- the context size which might be large if many results are returned
|
||||
- the query content which might include sensitive information
|
||||
- ...
|
||||
|
||||
## Performance Optimization
|
||||
|
||||
### Query Optimization
|
||||
|
||||
You may want to help the LLM formulate better queries by including something like this in the system prompt:
|
||||
|
||||
```
|
||||
When using web search:
|
||||
- Use specific, focused queries
|
||||
- Include relevant time periods if needed
|
||||
- Avoid unnecessary words
|
||||
- Combine related terms
|
||||
```
|
||||
|
||||
## See Also
|
||||
|
||||
- [Tools Overview](../tools.md)
|
||||
- [Brave Web Search Tool](web_search_brave.md)
|
||||
- [Web Search Configuration](../llm-configuration.md)
|
||||
- [Environment Variables](../env.md)
|
||||
|
||||
+1
-50
@@ -83,55 +83,6 @@ If you already have CRLF line endings in your local repository, the **best appro
|
||||
git commit -m "✏️(project) Fix line endings to LF"
|
||||
```
|
||||
|
||||
## Minio Permission Issues on Windows
|
||||
|
||||
### Problem Description
|
||||
|
||||
On Windows, you may encounter permission-related errors when running Minio in development mode with Docker Compose. This typically happens because:
|
||||
|
||||
- **Windows file permissions** don't map well to Unix-style user IDs used in Docker containers
|
||||
- **Docker Desktop** may have issues with user mapping when using the `DOCKER_USER` environment variable
|
||||
- **Minio container** fails to start or access volumes due to permission conflicts
|
||||
|
||||
### Common Symptoms
|
||||
|
||||
- Minio container fails to start with permission denied errors
|
||||
- Error messages related to file system permissions in Minio logs
|
||||
- Unable to create or access buckets in the development environment
|
||||
- Docker Compose showing Minio service as unhealthy or exited
|
||||
|
||||
### Solution for Windows Users
|
||||
|
||||
If you encounter Minio permission issues on Windows, you can temporarily disable user mapping for the Minio service:
|
||||
|
||||
1. **Open the `compose.yml` file**
|
||||
|
||||
2. **Comment out the user directive** in the `minio` service section:
|
||||
```yaml
|
||||
minio:
|
||||
# user: ${DOCKER_USER:-1000} # Comment this line on Windows if permission issues occur
|
||||
image: minio/minio
|
||||
environment:
|
||||
- MINIO_ROOT_USER=conversations
|
||||
- MINIO_ROOT_PASSWORD=password
|
||||
# ... rest of the configuration
|
||||
```
|
||||
|
||||
3. **Restart the services**:
|
||||
```bash
|
||||
make run
|
||||
```
|
||||
|
||||
### Why This Works
|
||||
|
||||
- Commenting out the `user` directive allows the Minio container to run with its default user
|
||||
- This bypasses Windows-specific permission mapping issues
|
||||
- The container will have the necessary permissions to access and manage the mounted volumes
|
||||
|
||||
### Note
|
||||
|
||||
This is a **development-only workaround**. In production environments, proper user mapping and security considerations should be maintained according to your deployment requirements.
|
||||
|
||||
## Frontend File Watching Issues on Windows
|
||||
|
||||
### Problem Description
|
||||
@@ -189,4 +140,4 @@ Add the `WATCHPACK_POLLING=true` environment variable to the frontend-developmen
|
||||
|
||||
### Note
|
||||
|
||||
This setting is primarily needed for Windows users. Linux and macOS users typically don't need this setting as file watching works correctly by default on those platforms.
|
||||
This setting is primarily needed for Windows users. Linux and macOS users typically don't need this setting as file watching works correctly by default on those platforms.
|
||||
|
||||
@@ -26,9 +26,6 @@ CONVERSATIONS_BASE_URL="http://localhost:8072"
|
||||
|
||||
# Media
|
||||
STORAGES_STATICFILES_BACKEND=django.contrib.staticfiles.storage.StaticFilesStorage
|
||||
AWS_S3_ENDPOINT_URL=http://minio:9000
|
||||
AWS_S3_ACCESS_KEY_ID=conversations
|
||||
AWS_S3_SECRET_ACCESS_KEY=password
|
||||
MEDIA_BASE_URL=http://localhost:8083
|
||||
|
||||
# OIDC
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
# For the CI job test-e2e
|
||||
BURST_THROTTLE_RATES="200/minute"
|
||||
DJANGO_SERVER_TO_SERVER_API_TOKENS=test-e2e
|
||||
SUSTAINED_THROTTLE_RATES="200/hour"
|
||||
|
||||
# LLM
|
||||
LLM_CONFIGURATION_FILE_PATH = /app/conversations/configuration/llm/default.e2e.json
|
||||
|
||||
# Features
|
||||
FEATURE_FLAG_WEB_SEARCH=ENABLED
|
||||
FEATURE_FLAG_DOCUMENT_UPLOAD=ENABLED
|
||||
|
||||
@@ -1,19 +0,0 @@
|
||||
FROM python:3.13.3-alpine
|
||||
|
||||
# Upgrade pip to its latest release to speed up dependencies installation
|
||||
RUN python -m pip install --upgrade pip setuptools lorem-text
|
||||
|
||||
# Upgrade system packages to install security updates
|
||||
RUN apk update && \
|
||||
apk upgrade
|
||||
|
||||
RUN apk add --no-cache git
|
||||
|
||||
# Install the package
|
||||
RUN pip install git+https://github.com/etalab-ia/openmockllm.git
|
||||
|
||||
# Expose the default port
|
||||
EXPOSE 8000
|
||||
|
||||
# Set default command
|
||||
CMD ["openmockllm", "--host", "0.0.0.0", "--port", "8000"]
|
||||
@@ -1,19 +0,0 @@
|
||||
[OpenMockLLM](https://github.com/etalab-ia/OpenMockLLM) is a FastAPI-based mock LLM API server that simulates
|
||||
several Large Language Model API providers.
|
||||
|
||||
This is a simple docker image to run the server for testing and development purposes (E2E tests mainly).
|
||||
|
||||
It's a bit overkill to have a dedicated image for that, but it allows simple E2E stack with docker-compose since
|
||||
our code is also run in Docker containers.
|
||||
|
||||
## Build and Run manually
|
||||
|
||||
```bash
|
||||
docker build -t openmockllm .
|
||||
docker run -p 8000:8000 openmockllm
|
||||
```
|
||||
|
||||
## Next steps
|
||||
|
||||
- Add more chat completion behaviors (specific text streaming, function calling, etc.)
|
||||
- Pin a specific OpenMockLLM version in the Dockerfile
|
||||
@@ -1,6 +1,5 @@
|
||||
"""Admin classes for activation codes application."""
|
||||
|
||||
from django.conf import settings
|
||||
from django.contrib import admin
|
||||
from django.utils.html import format_html, format_html_join
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
@@ -275,8 +274,6 @@ class UserRegistrationRequestAdmin(admin.ModelAdmin):
|
||||
|
||||
list_filter = ("created_at",)
|
||||
|
||||
actions = ["add_to_brevo_waiting_list", "remove_from_brevo_waiting_list"]
|
||||
|
||||
def user_display(self, obj):
|
||||
"""Display user's full name."""
|
||||
return obj.user.email or str(obj.user.pk)
|
||||
@@ -289,61 +286,3 @@ class UserRegistrationRequestAdmin(admin.ModelAdmin):
|
||||
|
||||
has_user_activation.boolean = True
|
||||
has_user_activation.short_description = _("Has used activation code")
|
||||
|
||||
@admin.action(description=_("Add selected users to Brevo waiting list"))
|
||||
def add_to_brevo_waiting_list(self, request, queryset):
|
||||
"""Add selected users to Brevo waiting list."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from core.brevo import add_user_to_brevo_list # noqa: PLC0415
|
||||
|
||||
registration_to_send = queryset.filter(
|
||||
user_activation__isnull=True,
|
||||
)
|
||||
|
||||
_total_emails = 0
|
||||
for i in range(0, registration_to_send.count(), 150):
|
||||
batch = registration_to_send[i : i + 150]
|
||||
emails = [reg.user.email for reg in batch if reg.user.email]
|
||||
if emails:
|
||||
add_user_to_brevo_list(emails, settings.BREVO_WAITING_LIST_ID)
|
||||
_total_emails += len(emails)
|
||||
|
||||
if _total_emails:
|
||||
self.message_user(
|
||||
request,
|
||||
_("Added %(count)d user(s) to Brevo waiting list.") % {"count": _total_emails},
|
||||
)
|
||||
else:
|
||||
self.message_user(
|
||||
request,
|
||||
_("No valid email address found in selected registrations."),
|
||||
level="warning",
|
||||
)
|
||||
|
||||
@admin.action(description=_("Remove selected users from Brevo waiting list"))
|
||||
def remove_from_brevo_waiting_list(self, request, queryset):
|
||||
"""Remove selected users from Brevo waiting list."""
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from core.brevo import remove_user_from_brevo_list # noqa: PLC0415
|
||||
|
||||
registration_to_send = queryset.filter(
|
||||
user_activation__isnull=False,
|
||||
)
|
||||
_total_emails = 0
|
||||
for i in range(0, registration_to_send.count(), 150):
|
||||
batch = registration_to_send[i : i + 150]
|
||||
emails = [reg.user.email for reg in batch if reg.user.email]
|
||||
if emails:
|
||||
remove_user_from_brevo_list(emails, settings.BREVO_WAITING_LIST_ID)
|
||||
_total_emails += len(emails)
|
||||
if _total_emails:
|
||||
self.message_user(
|
||||
request,
|
||||
_("Removed %(count)d user(s) from Brevo waiting list.") % {"count": _total_emails},
|
||||
)
|
||||
else:
|
||||
self.message_user(
|
||||
request,
|
||||
_("No valid email address found in selected registrations."),
|
||||
level="warning",
|
||||
)
|
||||
|
||||
@@ -6,14 +6,12 @@ import logging
|
||||
import secrets
|
||||
import string
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import ValidationError
|
||||
from django.core.validators import RegexValidator
|
||||
from django.db import IntegrityError, models, transaction
|
||||
from django.utils import timezone
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
|
||||
from core.brevo import add_user_to_brevo_list, remove_user_from_brevo_list
|
||||
from core.models import BaseModel, User
|
||||
|
||||
from activation_codes.exceptions import InvalidCodeError, UserAlreadyActivatedError
|
||||
@@ -136,24 +134,12 @@ class ActivationCode(BaseModel):
|
||||
_("You have already activated your account")
|
||||
) from exc
|
||||
|
||||
existing_registration = bool(
|
||||
UserRegistrationRequest.objects.filter(user=user).update(user_activation=activation)
|
||||
)
|
||||
if existing_registration:
|
||||
transaction.on_commit(
|
||||
lambda: remove_user_from_brevo_list(
|
||||
[user.email], settings.BREVO_WAITING_LIST_ID
|
||||
)
|
||||
)
|
||||
UserRegistrationRequest.objects.filter(user=user).update(user_activation=activation)
|
||||
|
||||
# Increment usage counter safely under the same lock.
|
||||
locked_code.current_uses += 1
|
||||
locked_code.save(update_fields=["current_uses", "updated_at"])
|
||||
|
||||
transaction.on_commit(
|
||||
lambda: add_user_to_brevo_list([user.email], settings.BREVO_FOLLOWUP_LIST_ID)
|
||||
)
|
||||
|
||||
if locked_code.max_uses > 0 and locked_code.current_uses >= locked_code.max_uses:
|
||||
logger.warning("Activation code %s has reached its maximum uses", locked_code.code)
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
"""Tests for activation_codes models."""
|
||||
|
||||
import json
|
||||
from datetime import timedelta
|
||||
|
||||
from django.core.exceptions import ValidationError
|
||||
@@ -8,18 +7,12 @@ from django.db.models import ProtectedError
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
|
||||
from core.factories import UserFactory
|
||||
|
||||
from activation_codes.exceptions import InvalidCodeError, UserAlreadyActivatedError
|
||||
from activation_codes.factories import ActivationCodeFactory, UserActivationFactory
|
||||
from activation_codes.models import (
|
||||
ActivationCode,
|
||||
UserActivation,
|
||||
UserRegistrationRequest,
|
||||
generate_activation_code,
|
||||
)
|
||||
from activation_codes.models import ActivationCode, UserActivation, generate_activation_code
|
||||
|
||||
|
||||
@pytest.mark.django_db
|
||||
@@ -277,52 +270,3 @@ def test_user_activation_ordering():
|
||||
|
||||
activations = list(UserActivation.objects.all())
|
||||
assert activations == [activation2, activation1]
|
||||
|
||||
|
||||
@responses.activate
|
||||
@pytest.mark.django_db(transaction=True)
|
||||
def test_activation_code_use_success_notify_brevo(settings):
|
||||
"""Test successfully using an activation code and notify Brevo."""
|
||||
settings.BREVO_API_KEY = "test_brevo_api_key"
|
||||
settings.BREVO_WAITING_LIST_ID = "test_waiting_list_id"
|
||||
settings.BREVO_FOLLOWUP_LIST_ID = "test_followup_list_name"
|
||||
|
||||
brevo_remove_mock = responses.post(
|
||||
"https://api.brevo.com/v3/contacts/lists/test_waiting_list_id/contacts/remove",
|
||||
json={"message": "Contacts added successfully"},
|
||||
status=201,
|
||||
)
|
||||
|
||||
brevo_create_contact = responses.post(
|
||||
"https://api.brevo.com/v3/contacts",
|
||||
status=200,
|
||||
)
|
||||
|
||||
brevo_add_mock = responses.post(
|
||||
"https://api.brevo.com/v3/contacts/lists/test_followup_list_name/contacts/add",
|
||||
json={"message": "Contacts added successfully"},
|
||||
status=201,
|
||||
)
|
||||
|
||||
user = UserFactory()
|
||||
registration = UserRegistrationRequest.objects.create(user=user)
|
||||
activation_code = ActivationCodeFactory()
|
||||
activation = activation_code.use(user)
|
||||
|
||||
registration.refresh_from_db()
|
||||
assert registration.user_activation == activation
|
||||
|
||||
assert len(brevo_remove_mock.calls) == 1
|
||||
assert brevo_remove_mock.calls[0].request.headers["api-key"] == "test_brevo_api_key"
|
||||
assert json.loads(brevo_remove_mock.calls[0].request.body) == {"emails": [user.email]}
|
||||
|
||||
assert len(brevo_create_contact.calls) == 1
|
||||
assert brevo_create_contact.calls[0].request.headers["api-key"] == "test_brevo_api_key"
|
||||
assert json.loads(brevo_create_contact.calls[0].request.body) == {
|
||||
"email": user.email,
|
||||
"updateEnabled": True,
|
||||
}
|
||||
|
||||
assert len(brevo_add_mock.calls) == 1
|
||||
assert brevo_add_mock.calls[0].request.headers["api-key"] == "test_brevo_api_key"
|
||||
assert json.loads(brevo_add_mock.calls[0].request.body) == {"emails": [user.email]}
|
||||
|
||||
@@ -1,13 +1,11 @@
|
||||
"""Tests for activation_codes viewsets."""
|
||||
|
||||
import json
|
||||
from datetime import timedelta
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
from rest_framework import status
|
||||
|
||||
from core.factories import UserFactory
|
||||
@@ -323,89 +321,3 @@ def test_validate_code_registered_user(api_client):
|
||||
|
||||
_registration.refresh_from_db()
|
||||
assert _registration.user_activation.activation_code == activation_code
|
||||
|
||||
|
||||
@responses.activate
|
||||
@pytest.mark.django_db
|
||||
def test_register_email_success_brevo(api_client, settings):
|
||||
"""Test successfully registering an email and notify Brevo."""
|
||||
settings.BREVO_API_KEY = "test_brevo_api_key"
|
||||
settings.BREVO_WAITING_LIST_ID = "test_waiting_list_id"
|
||||
|
||||
brevo_create_contact = responses.post(
|
||||
"https://api.brevo.com/v3/contacts",
|
||||
status=200,
|
||||
)
|
||||
|
||||
brevo_mock = responses.post(
|
||||
"https://api.brevo.com/v3/contacts/lists/test_waiting_list_id/contacts/add",
|
||||
json={"message": "Contacts added successfully"},
|
||||
status=201,
|
||||
)
|
||||
|
||||
user = UserFactory()
|
||||
api_client.force_authenticate(user=user)
|
||||
|
||||
response = api_client.post(
|
||||
"/api/v1.0/activation/register/",
|
||||
{},
|
||||
)
|
||||
assert response.status_code == status.HTTP_201_CREATED
|
||||
assert response.data["code"] == "registration-successful"
|
||||
|
||||
registration = UserRegistrationRequest.objects.get(user=user)
|
||||
assert registration.user == user
|
||||
|
||||
assert len(brevo_create_contact.calls) == 1
|
||||
assert brevo_create_contact.calls[0].request.headers["api-key"] == "test_brevo_api_key"
|
||||
assert json.loads(brevo_create_contact.calls[0].request.body) == {
|
||||
"email": user.email,
|
||||
"updateEnabled": True,
|
||||
}
|
||||
|
||||
assert len(brevo_mock.calls) == 1
|
||||
assert brevo_mock.calls[0].request.headers["api-key"] == "test_brevo_api_key"
|
||||
assert json.loads(brevo_mock.calls[0].request.body) == {"emails": [user.email]}
|
||||
|
||||
# Register again to test idempotency
|
||||
response = api_client.post(
|
||||
"/api/v1.0/activation/register/",
|
||||
{},
|
||||
)
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.data["code"] == "registration-successful"
|
||||
|
||||
assert len(brevo_mock.calls) == 1 # No new call made
|
||||
|
||||
|
||||
@responses.activate
|
||||
@pytest.mark.django_db
|
||||
def test_register_email_success_brevo_fails(api_client, settings):
|
||||
"""Test successfully registering an email, even if Brevo fails."""
|
||||
settings.BREVO_API_KEY = "test_brevo_api_key"
|
||||
settings.BREVO_WAITING_LIST_ID = "test_waiting_list_id"
|
||||
|
||||
_brevo_create_contact = responses.post(
|
||||
"https://api.brevo.com/v3/contacts",
|
||||
status=200,
|
||||
)
|
||||
|
||||
brevo_mock = responses.post(
|
||||
"https://api.brevo.com/v3/contacts/lists/test_waiting_list_id/contacts/add",
|
||||
status=400,
|
||||
)
|
||||
|
||||
user = UserFactory()
|
||||
api_client.force_authenticate(user=user)
|
||||
|
||||
response = api_client.post(
|
||||
"/api/v1.0/activation/register/",
|
||||
{},
|
||||
)
|
||||
assert response.status_code == status.HTTP_201_CREATED
|
||||
assert response.data["code"] == "registration-successful"
|
||||
|
||||
registration = UserRegistrationRequest.objects.get(user=user)
|
||||
assert registration.user == user
|
||||
|
||||
assert len(brevo_mock.calls) == 1
|
||||
|
||||
@@ -10,7 +10,6 @@ from rest_framework import status, viewsets
|
||||
from rest_framework.decorators import action
|
||||
from rest_framework.response import Response
|
||||
|
||||
from core.brevo import add_user_to_brevo_list
|
||||
from core.permissions import IsAuthenticated
|
||||
|
||||
from . import models, serializers
|
||||
@@ -138,10 +137,6 @@ class ActivationViewSet(viewsets.GenericViewSet):
|
||||
status=status.HTTP_200_OK,
|
||||
)
|
||||
|
||||
add_user_to_brevo_list(
|
||||
[serializer.validated_data["user"].email], settings.BREVO_WAITING_LIST_ID
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Registered email %s for activation notifications",
|
||||
serializer.validated_data["user"].email,
|
||||
|
||||
@@ -3,12 +3,11 @@
|
||||
import json
|
||||
import logging
|
||||
from io import BytesIO
|
||||
from typing import List, Optional
|
||||
from typing import Optional
|
||||
from urllib.parse import urljoin
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
import httpx
|
||||
import requests
|
||||
|
||||
from chat.agent_rag.albert_api_constants import Searches
|
||||
@@ -33,13 +32,9 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
- Perform a search operation using the Albert API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_id: Optional[str] = None,
|
||||
read_only_collection_id: Optional[List[str]] = None,
|
||||
):
|
||||
def __init__(self, collection_id: Optional[str] = None):
|
||||
# Initialize any necessary parameters or configurations here
|
||||
super().__init__(collection_id, read_only_collection_id)
|
||||
super().__init__(collection_id)
|
||||
self._base_url = settings.ALBERT_API_URL
|
||||
self._headers = {
|
||||
"Authorization": f"Bearer {settings.ALBERT_API_KEY}",
|
||||
@@ -70,27 +65,6 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
self.collection_id = str(response.json()["id"])
|
||||
return self.collection_id
|
||||
|
||||
async def acreate_collection(self, name: str, description: Optional[str] = None) -> str:
|
||||
"""
|
||||
Create a temporary collection for the search operation.
|
||||
This method should handle the logic to create or retrieve an existing collection.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
self._collections_endpoint,
|
||||
headers=self._headers,
|
||||
json={
|
||||
"name": name,
|
||||
"description": description or self._default_collection_description,
|
||||
"visibility": "private",
|
||||
},
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
self.collection_id = str(response.json()["id"])
|
||||
return self.collection_id
|
||||
|
||||
def delete_collection(self) -> None:
|
||||
"""
|
||||
Delete the current collection
|
||||
@@ -102,18 +76,6 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
async def adelete_collection(self) -> None:
|
||||
"""
|
||||
Asynchronously delete the current collection
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
|
||||
response = await client.delete(
|
||||
urljoin(f"{self._collections_endpoint}/", self.collection_id),
|
||||
headers=self._headers,
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
def parse_pdf_document(self, name: str, content_type: str, content: BytesIO) -> str:
|
||||
"""
|
||||
Parse the PDF document content and return the text content.
|
||||
@@ -188,52 +150,23 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
logger.debug(response.json())
|
||||
response.raise_for_status()
|
||||
|
||||
async def astore_document(self, name: str, content: str) -> None:
|
||||
"""
|
||||
Store the document content in the Albert collection.
|
||||
This method should handle the logic to send the document content to the Albert API.
|
||||
|
||||
Args:
|
||||
name (str): The name of the document.
|
||||
content (str): The content of the document in Markdown format.
|
||||
"""
|
||||
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
urljoin(self._base_url, self._documents_endpoint),
|
||||
headers=self._headers,
|
||||
files={
|
||||
"file": (f"{name}.md", BytesIO(content.encode("utf-8")), "text/markdown"),
|
||||
},
|
||||
data={
|
||||
"collection": int(self.collection_id),
|
||||
"metadata": json.dumps({"document_name": name}), # undocumented API
|
||||
},
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
logger.debug(response.json())
|
||||
response.raise_for_status()
|
||||
|
||||
def search(self, query, results_count: int = 4) -> RAGWebResults:
|
||||
def search(self, query) -> RAGWebResults:
|
||||
"""
|
||||
Perform a search using the Albert API based on the provided query.
|
||||
|
||||
Args:
|
||||
query (str): The search query.
|
||||
results_count (int): The number of results to return.
|
||||
|
||||
Returns:
|
||||
RAGWebResults: The search results.
|
||||
"""
|
||||
collection_ids = self.get_all_collection_ids() # might raise RuntimeError
|
||||
|
||||
response = requests.post(
|
||||
urljoin(self._base_url, self._search_endpoint),
|
||||
headers=self._headers,
|
||||
json={
|
||||
"collections": collection_ids,
|
||||
"collections": [int(self.collection_id)],
|
||||
"prompt": query,
|
||||
"score_threshold": 0.6,
|
||||
"k": results_count, # Number of chunks to return from the search
|
||||
},
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
@@ -255,50 +188,3 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
completion_tokens=searches.usage.completion_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
async def asearch(self, query, results_count: int = 4) -> RAGWebResults:
|
||||
"""
|
||||
Perform an asynchronous search using the Albert API based on the provided query.
|
||||
|
||||
Args:
|
||||
query (str): The search query.
|
||||
results_count (int): The number of results to return.
|
||||
|
||||
Returns:
|
||||
RAGWebResults: The search results.
|
||||
"""
|
||||
collection_ids = self.get_all_collection_ids() # might raise RuntimeError
|
||||
|
||||
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
|
||||
response = await client.post(
|
||||
urljoin(self._base_url, self._search_endpoint),
|
||||
headers=self._headers,
|
||||
json={
|
||||
"collections": collection_ids,
|
||||
"prompt": query,
|
||||
"score_threshold": 0.6,
|
||||
"k": results_count, # Number of chunks to return from the search
|
||||
},
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
|
||||
logger.debug("Search response: %s %s", response.text, response.status_code)
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
searches = Searches(**response.json())
|
||||
|
||||
return RAGWebResults(
|
||||
data=[
|
||||
RAGWebResult(
|
||||
url=result.chunk.metadata["document_name"],
|
||||
content=result.chunk.content,
|
||||
score=result.score,
|
||||
)
|
||||
for result in searches.data
|
||||
],
|
||||
usage=RAGWebUsage(
|
||||
prompt_tokens=searches.usage.prompt_tokens,
|
||||
completion_tokens=searches.usage.completion_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -1,11 +1,9 @@
|
||||
"""Implementation of the Albert API for RAG document search."""
|
||||
|
||||
import logging
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from contextlib import contextmanager
|
||||
from io import BytesIO
|
||||
from typing import List, Optional
|
||||
|
||||
from asgiref.sync import sync_to_async
|
||||
from typing import Optional
|
||||
|
||||
from chat.agent_rag.constants import RAGWebResults
|
||||
|
||||
@@ -15,51 +13,11 @@ logger = logging.getLogger(__name__)
|
||||
class BaseRagBackend:
|
||||
"""Base class for RAG backends."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
collection_id: Optional[str] = None,
|
||||
read_only_collection_id: Optional[List[str]] = None,
|
||||
):
|
||||
"""
|
||||
Backend settings.
|
||||
|
||||
Collection ID is required for RAG operations, where you want to manage the collection
|
||||
lifecycle (create/delete).
|
||||
Read-only collection IDs can be used to access existing collections
|
||||
without managing their lifecycle.
|
||||
|
||||
Collection ID and read-only collection IDs are separated in the implementation to prevent
|
||||
unwanted actions.
|
||||
|
||||
Args:
|
||||
collection_id (Optional[str]): The collection ID for managing the collection lifecycle.
|
||||
read_only_collection_id (Optional[List[str]]): List of read-only collection IDs.
|
||||
"""
|
||||
def __init__(self, collection_id: Optional[str] = None):
|
||||
"""Backend settings."""
|
||||
self.collection_id = collection_id
|
||||
self.read_only_collection_id = read_only_collection_id or []
|
||||
self._default_collection_description = "Temporary collection for RAG document search"
|
||||
|
||||
def get_all_collection_ids(self) -> List[str]:
|
||||
"""
|
||||
Get all collection IDs, including the main collection ID and read-only collection IDs.
|
||||
|
||||
Returns:
|
||||
List[str]: List of all collection IDs.
|
||||
Raises:
|
||||
RuntimeError: If neither collection_id nor read_only_collection_id is provided.
|
||||
"""
|
||||
if not self.collection_id and not self.read_only_collection_id:
|
||||
raise RuntimeError("The RAG backend requires collection_id or read_only_collection_id")
|
||||
|
||||
collection_ids = []
|
||||
if self.collection_id:
|
||||
collection_ids.append(int(self.collection_id))
|
||||
if self.read_only_collection_id:
|
||||
collection_ids.extend(
|
||||
[int(collection_id) for collection_id in self.read_only_collection_id]
|
||||
)
|
||||
return collection_ids
|
||||
|
||||
def create_collection(self, name: str, description: Optional[str] = None) -> str:
|
||||
"""
|
||||
Create a temporary collection for the search operation.
|
||||
@@ -67,13 +25,6 @@ class BaseRagBackend:
|
||||
"""
|
||||
raise NotImplementedError("Must be implemented in subclass.")
|
||||
|
||||
async def acreate_collection(self, name: str, description: Optional[str] = None) -> str:
|
||||
"""
|
||||
Create a temporary collection for the search operation.
|
||||
This method should handle the logic to create or retrieve an existing collection.
|
||||
"""
|
||||
return await sync_to_async(self.create_collection)(name=name, description=description)
|
||||
|
||||
def parse_document(self, name: str, content_type: str, content: BytesIO):
|
||||
"""
|
||||
Parse the document and prepare it for the search operation.
|
||||
@@ -92,8 +43,8 @@ class BaseRagBackend:
|
||||
|
||||
def store_document(self, name: str, content: str) -> None:
|
||||
"""
|
||||
Store the document content in the collection.
|
||||
This method should handle the logic to send the document content to the API.
|
||||
Store the document content in the Albert collection.
|
||||
This method should handle the logic to send the document content to the Albert API.
|
||||
|
||||
Args:
|
||||
name (str): The name of the document.
|
||||
@@ -101,17 +52,6 @@ class BaseRagBackend:
|
||||
"""
|
||||
raise NotImplementedError("Must be implemented in subclass.")
|
||||
|
||||
async def astore_document(self, name: str, content: str) -> None:
|
||||
"""
|
||||
Store the document content in the collection.
|
||||
This method should handle the logic to send the document content to the API.
|
||||
|
||||
Args:
|
||||
name (str): The name of the document.
|
||||
content (str): The content of the document in Markdown format.
|
||||
"""
|
||||
return await sync_to_async(self.store_document)(name=name, content=content)
|
||||
|
||||
def parse_and_store_document(self, name: str, content_type: str, content: BytesIO) -> str:
|
||||
"""
|
||||
Parse the document and store it in the Albert collection.
|
||||
@@ -135,25 +75,12 @@ class BaseRagBackend:
|
||||
"""
|
||||
raise NotImplementedError("Must be implemented in subclass.")
|
||||
|
||||
async def adelete_collection(self) -> None:
|
||||
"""
|
||||
Delete the collection.
|
||||
This method should handle the logic to delete the collection from the backend.
|
||||
"""
|
||||
return await sync_to_async(self.delete_collection)()
|
||||
|
||||
def search(self, query, results_count: int = 4) -> RAGWebResults:
|
||||
def search(self, query) -> RAGWebResults:
|
||||
"""
|
||||
Search the collection for the given query.
|
||||
"""
|
||||
raise NotImplementedError("Must be implemented in subclass.")
|
||||
|
||||
async def asearch(self, query, results_count: int = 4) -> RAGWebResults:
|
||||
"""
|
||||
Search the collection for the given query.
|
||||
"""
|
||||
return await sync_to_async(self.search)(query=query, results_count=results_count)
|
||||
|
||||
@classmethod
|
||||
@contextmanager
|
||||
def temporary_collection(cls, name: str, description: Optional[str] = None):
|
||||
@@ -165,15 +92,3 @@ class BaseRagBackend:
|
||||
yield backend
|
||||
finally:
|
||||
backend.delete_collection()
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def temporary_collection_async(cls, name: str, description: Optional[str] = None):
|
||||
"""Context manager for RAG backend with temporary collections."""
|
||||
backend = cls()
|
||||
|
||||
await backend.acreate_collection(name=name, description=description)
|
||||
try:
|
||||
yield backend
|
||||
finally:
|
||||
await backend.adelete_collection()
|
||||
|
||||
@@ -169,12 +169,11 @@ class AlbertRagDocumentSearch:
|
||||
self._store_document(name, document_content)
|
||||
return document_content
|
||||
|
||||
def search(self, query, results_count: int = 4) -> RAGWebResults:
|
||||
def search(self, query):
|
||||
"""
|
||||
Perform a search using the Albert API based on the provided query.
|
||||
|
||||
:param query: The search query string.
|
||||
:param results_count: The number of results to return.
|
||||
:return: Search results from the Albert API.
|
||||
"""
|
||||
response = requests.post(
|
||||
@@ -184,7 +183,6 @@ class AlbertRagDocumentSearch:
|
||||
"collections": [self.collection_id],
|
||||
"prompt": query,
|
||||
"score_threshold": 0.6,
|
||||
"k": results_count, # Number of chunks to return from the search
|
||||
},
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
|
||||
@@ -1,15 +1,10 @@
|
||||
"""Build the main conversation agent."""
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.utils import formats, timezone
|
||||
|
||||
from pydantic_ai import ModelMessage
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
|
||||
from core.enums import get_language_name
|
||||
|
||||
from .base import BaseAgent
|
||||
@@ -17,79 +12,6 @@ from .base import BaseAgent
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
MOCKED_RESPONSE = """
|
||||
# **Ode to the AI Assistant** 🤖✨
|
||||
|
||||
In Paris streets where old meets new, 🗼🇫🇷
|
||||
A helper bright in digital hue,
|
||||
With circuits fast and code so tight,
|
||||
The LaSuite’s bot—oh, what a sight! 🌟
|
||||
|
||||
**A chatbot kind**, with wittiness so grand, 💬💡
|
||||
It lends a hand to all the land,
|
||||
From civil servants, bold and wise,
|
||||
To those who seek with hopeful eyes.
|
||||
|
||||
It answers quick, it never tires, ⚡🔄
|
||||
With facts and tips to quench desires,
|
||||
A guide so keen, a friend so true,
|
||||
It’s there for **you**—yes, me and you!
|
||||
|
||||
With **Markdown flair** and emoji cheer, 📝🎨
|
||||
It makes the complex crystal clear,
|
||||
From drafts to code, from sums to prose,
|
||||
It helps the knowledge overflow!
|
||||
|
||||
Oh, **DINUM’s gem**, so sharp, so bright, 💎🌐
|
||||
A beacon in the tech’s vast night,
|
||||
It crafts, it checks, it summarizes,
|
||||
With grace that never compromises.
|
||||
|
||||
It **summarizes** the long, the deep, 📚🔍
|
||||
So secrets no more need to sleep,
|
||||
It finds the gems in data’s sea,
|
||||
And sets the truth right there—**for free!**
|
||||
|
||||
It **corrects mistakes** with gentle art, ✍️🔄
|
||||
It soothes the mind, it warms the heart,
|
||||
No judgment cast, no frown, no sigh,
|
||||
Just help that’s always standing by.
|
||||
|
||||
It **generates code** with swift command, 💻🔥
|
||||
A developer’s dream, first-hand,
|
||||
From Python lines to scripts so neat,
|
||||
It turns the tough to *sweet* and *sweet*!
|
||||
|
||||
It **brainstorms ideas**, bold and new, 🧠💡
|
||||
It paints the sky in every hue,
|
||||
From plans to dreams, from start to end,
|
||||
It’s more than code—it’s **trend**, it’s **friend**!
|
||||
|
||||
So here’s to you, **Assistant’s pride**, 🏆🎉
|
||||
The bot that’s always by our side,
|
||||
With every prompt, with every line,
|
||||
You make our digital world **divine**!
|
||||
|
||||
May you keep shining, bright and true, 🌟🤖
|
||||
The helper every team should woo,
|
||||
For in this age of bits and bytes,
|
||||
You’re **human touch** in tech’s bright lights!
|
||||
"""
|
||||
|
||||
|
||||
async def mocked_agent_model(_messages: list[ModelMessage], _info: AgentInfo):
|
||||
"""
|
||||
Mocked agent model for testing purposes on deployed instances.
|
||||
|
||||
This one only fakes a streamed responses. We could also fake tool calls later.
|
||||
"""
|
||||
|
||||
yield "Here is a mocked response (no LLM called)\n---\n"
|
||||
for i in range(0, len(MOCKED_RESPONSE), 4):
|
||||
yield MOCKED_RESPONSE[i : i + 4]
|
||||
await asyncio.sleep(0.03)
|
||||
|
||||
|
||||
@dataclasses.dataclass(init=False)
|
||||
class ConversationAgent(BaseAgent):
|
||||
"""Conversation agent with custom behavior."""
|
||||
@@ -98,10 +20,6 @@ class ConversationAgent(BaseAgent):
|
||||
"""Initialize the conversation agent."""
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Do not call the real model on deployed instances if the setting is enabled
|
||||
if settings.WARNING_MOCK_CONVERSATION_AGENT:
|
||||
self._model = FunctionModel(stream_function=mocked_agent_model)
|
||||
|
||||
@self.system_prompt
|
||||
def add_the_date() -> str:
|
||||
"""
|
||||
|
||||
@@ -1,94 +0,0 @@
|
||||
"""
|
||||
ImageUrl processors and utilities.
|
||||
|
||||
Allow to manage local image URLs in messages, replacing them with presigned S3 URLs
|
||||
for the LLM to access them, and then reverting them back to local URLs when
|
||||
storing the messages in the database.
|
||||
"""
|
||||
|
||||
from typing import Dict, Iterable
|
||||
|
||||
from pydantic_ai import DocumentUrl, ImageUrl, ModelMessage, ModelRequest, UserPromptPart
|
||||
|
||||
from core.file_upload.utils import generate_retrieve_policy
|
||||
|
||||
from chat.models import ChatConversation
|
||||
|
||||
|
||||
def update_local_urls(
|
||||
conversation: ChatConversation,
|
||||
contents: Iterable[ImageUrl | DocumentUrl],
|
||||
updated_url: Dict[str, str] | None = None,
|
||||
) -> Iterable[ImageUrl | DocumentUrl]:
|
||||
"""
|
||||
Replace local image or document URLs in the content list to use presigned S3 URLs.
|
||||
⚠️Be careful, `media_contents` are replaced in place.
|
||||
|
||||
Args:
|
||||
conversation (ChatConversation): The chat conversation object.
|
||||
contents (Iterable[ImageUrl | DocumentUrl]): Iterable of UserContent objects.
|
||||
updated_url (Dict[str, str], optional): Dictionary to store
|
||||
mapping of original URLs to updated URLs.
|
||||
Returns:
|
||||
Iterable[ImageUrl | DocumentUrl]: Updated iterable of UserContent objects
|
||||
with presigned URLs.
|
||||
"""
|
||||
# When images are stored locally, there is no host in the URL, so we can
|
||||
# just check if the URL starts, frontend adds a prefix `/media-key/` to the key.
|
||||
local_media_url_prefix = "/media-key/"
|
||||
local_media_url_prefix_len = len(local_media_url_prefix)
|
||||
|
||||
# Filter only ImageUrl contents
|
||||
media_contents = (c for c in contents if isinstance(c, (ImageUrl, DocumentUrl)))
|
||||
|
||||
# Replace URLs with presigned URLs
|
||||
for content in media_contents:
|
||||
idx = content.url.find(local_media_url_prefix)
|
||||
|
||||
if idx == 0:
|
||||
_initial_url = str(content.url)
|
||||
key = content.url[local_media_url_prefix_len:]
|
||||
|
||||
# Security check: ensure the image belongs to the conversation, if yes,
|
||||
# the user had access to the endpoint, so they have access to the image.
|
||||
if not key.startswith(f"{conversation.pk}/"):
|
||||
# The LLM will throw an error when trying to access the image,
|
||||
# this is not perfect, but this should never happen in practice,
|
||||
# except if the user tampers with the conversation.
|
||||
continue
|
||||
|
||||
content.url = generate_retrieve_policy(key)
|
||||
if updated_url is not None:
|
||||
updated_url[content.url] = _initial_url
|
||||
|
||||
return contents
|
||||
|
||||
|
||||
def update_history_local_urls(
|
||||
conversation: ChatConversation, messages: list[ModelMessage]
|
||||
) -> list[ModelMessage]:
|
||||
"""
|
||||
Replace local image/documents URLs in the message list to use presigned S3 URLs.
|
||||
|
||||
⚠️Be careful, `messages` are replaced in place.
|
||||
|
||||
We don't need to store the mapping of updated URLs to original URLs here because
|
||||
this function is used when sending the history to the LLM (which is already stored
|
||||
in the database with local URLs).
|
||||
|
||||
Args:
|
||||
messages (list[ModelMessage]): List of ModelMessage objects.
|
||||
Returns:
|
||||
list[ModelMessage]: Updated list of ModelMessage objects with presigned URLs.
|
||||
"""
|
||||
# Filter only ModelRequest messages
|
||||
requests = (msg for msg in messages if isinstance(msg, ModelRequest))
|
||||
|
||||
for message in requests:
|
||||
# Filter only UserPromptPart parts
|
||||
user_parts = (part for part in message.parts if isinstance(part, UserPromptPart))
|
||||
|
||||
for part in user_parts:
|
||||
update_local_urls(conversation, part.content)
|
||||
|
||||
return messages
|
||||
@@ -5,6 +5,10 @@ import logging
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from ..models import ChatConversationContextKind
|
||||
from .base import BaseAgent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -21,3 +25,50 @@ class SummarizationAgent(BaseAgent):
|
||||
output_type=str,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
async def hand_off_to_summarization_agent(ctx: RunContext) -> ToolReturn:
|
||||
"""
|
||||
Summarize the documents for the user, only when asked for,
|
||||
the documents are in my context.
|
||||
"""
|
||||
summarization_agent = SummarizationAgent()
|
||||
|
||||
prompt = (
|
||||
"Do not mention the user request in your answer.\n"
|
||||
"User request:\n"
|
||||
"{user_prompt}\n\n"
|
||||
"Document contents:\n"
|
||||
"{documents_prompt}\n"
|
||||
)
|
||||
documents = ctx.deps.conversation.contexts.filter(
|
||||
kind=ChatConversationContextKind.DOCUMENT.value,
|
||||
)
|
||||
documents_prompt = "\n\n".join(
|
||||
[
|
||||
(
|
||||
"<document>\n"
|
||||
f"<name>\n{doc.name}\n</name>\n"
|
||||
f"<content>\n{doc.content}\n</content>\n"
|
||||
"</document>"
|
||||
)
|
||||
async for doc in documents
|
||||
]
|
||||
)
|
||||
|
||||
formatted_prompt = prompt.format(
|
||||
user_prompt=ctx.prompt,
|
||||
documents_prompt=documents_prompt,
|
||||
)
|
||||
|
||||
logger.debug("Summarize prompt: %s", formatted_prompt)
|
||||
|
||||
response = await summarization_agent.run(formatted_prompt, usage=ctx.usage)
|
||||
|
||||
logger.debug("Summarize response: %s", response)
|
||||
|
||||
return ToolReturn(
|
||||
return_value=response.output,
|
||||
content=response.output,
|
||||
metadata={"sources": {doc.name async for doc in documents}},
|
||||
)
|
||||
|
||||
@@ -7,32 +7,26 @@ changes are needed in views.py or tests.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack, ExitStack
|
||||
from io import BytesIO
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
from django.conf import settings
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.core.cache import cache
|
||||
from django.core.exceptions import ImproperlyConfigured
|
||||
from django.core.files.storage import default_storage
|
||||
from django.db.models import Q
|
||||
from django.utils.module_loading import import_string
|
||||
|
||||
from asgiref.sync import sync_to_async
|
||||
from langfuse import get_client
|
||||
from pydantic_ai import Agent, InstrumentationSettings, RunContext
|
||||
from pydantic_ai import Agent, ToolOutput
|
||||
from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
DocumentUrl,
|
||||
FunctionToolCallEvent,
|
||||
FunctionToolResultEvent,
|
||||
ImageUrl,
|
||||
ModelMessage,
|
||||
ModelMessagesTypeAdapter,
|
||||
ModelRequest,
|
||||
@@ -48,19 +42,14 @@ from pydantic_ai.messages import (
|
||||
ToolCallPartDelta,
|
||||
ToolReturn,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.result import FinalResult
|
||||
|
||||
from core.feature_flags.helpers import is_feature_enabled
|
||||
from core.file_upload.utils import generate_retrieve_policy
|
||||
|
||||
from chat import models
|
||||
from chat.agents.conversation import ConversationAgent
|
||||
from chat.agents.local_media_url_processors import (
|
||||
update_history_local_urls,
|
||||
update_local_urls,
|
||||
)
|
||||
from chat.agents.summarize import SummarizationAgent
|
||||
from chat.agents.summarize import hand_off_to_summarization_agent
|
||||
from chat.ai_sdk_types import (
|
||||
LanguageModelV1Source,
|
||||
SourceUIPart,
|
||||
@@ -73,9 +62,7 @@ from chat.clients.pydantic_ui_message_converter import (
|
||||
ui_message_to_user_content,
|
||||
)
|
||||
from chat.mcp_servers import get_mcp_servers
|
||||
from chat.tools.document_generic_search_rag import add_document_rag_search_tool_from_setting
|
||||
from chat.tools.document_search_rag import add_document_rag_search_tool
|
||||
from chat.tools.document_summarize import document_summarize
|
||||
from chat.vercel_ai_sdk.core import events_v4, events_v5
|
||||
from chat.vercel_ai_sdk.encoder import EventEncoder
|
||||
|
||||
@@ -118,8 +105,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
self.language = language # might be None
|
||||
self._last_stop_check = 0
|
||||
|
||||
self._langfuse_available = settings.LANGFUSE_ENABLED
|
||||
self._store_analytics = self._langfuse_available and user.allow_conversation_analytics
|
||||
self._store_analytics = settings.LANGFUSE_ENABLED and user.allow_conversation_analytics
|
||||
self.event_encoder = EventEncoder("v4") # Always use v4 for now
|
||||
|
||||
self._support_streaming = True
|
||||
@@ -140,15 +126,9 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
self.conversation_agent = ConversationAgent(
|
||||
model_hrid=self.model_hrid,
|
||||
language=self.language,
|
||||
instrument=InstrumentationSettings(
|
||||
include_binary_content=self._store_analytics,
|
||||
include_content=self._store_analytics,
|
||||
)
|
||||
if self._langfuse_available
|
||||
else False,
|
||||
instrument=self._store_analytics,
|
||||
deps_type=ContextDeps,
|
||||
)
|
||||
add_document_rag_search_tool_from_setting(self.conversation_agent, self.user)
|
||||
|
||||
@property
|
||||
def _stop_cache_key(self):
|
||||
@@ -183,7 +163,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
"""Return only the assistant text deltas (legacy text mode)."""
|
||||
await self._clean()
|
||||
with ExitStack() as stack:
|
||||
if self._langfuse_available:
|
||||
if self._store_analytics:
|
||||
span = stack.enter_context(get_client().start_as_current_span(name="conversation"))
|
||||
span.update_trace(user_id=str(self.user.sub), session_id=str(self.conversation.pk))
|
||||
|
||||
@@ -195,7 +175,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
"""Return Vercel-AI-SDK formatted events."""
|
||||
await self._clean()
|
||||
with ExitStack() as stack:
|
||||
if self._langfuse_available:
|
||||
if self._store_analytics:
|
||||
span = stack.enter_context(get_client().start_as_current_span(name="conversation"))
|
||||
span.update_trace(user_id=str(self.user.sub), session_id=str(self.conversation.pk))
|
||||
async for event in self._run_agent(messages, force_web_search):
|
||||
@@ -232,24 +212,10 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
# --------------------------------------------------------------------- #
|
||||
# Core agent runner
|
||||
# --------------------------------------------------------------------- #
|
||||
async def parse_input_documents(self, documents: List[BinaryContent | DocumentUrl]):
|
||||
async def parse_input_documents(self, documents: List[BinaryContent]):
|
||||
"""
|
||||
Parse and store input documents in the conversation's document store.
|
||||
"""
|
||||
# Early external document URL rejection
|
||||
if any(
|
||||
not document.url.startswith("/media-key/")
|
||||
for document in documents
|
||||
if isinstance(document, DocumentUrl)
|
||||
):
|
||||
raise ValueError("External document URL are not accepted yet.")
|
||||
if any(
|
||||
not document.url.startswith(f"/media-key/{self.conversation.pk}/")
|
||||
for document in documents
|
||||
if isinstance(document, DocumentUrl)
|
||||
):
|
||||
raise ValueError("Document URL does not belong to the conversation.")
|
||||
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
|
||||
document_store = document_store_backend(self.conversation.collection_id)
|
||||
@@ -262,48 +228,21 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
await self.conversation.asave(update_fields=["collection_id", "updated_at"])
|
||||
|
||||
for document in documents:
|
||||
key = None
|
||||
if isinstance(document, DocumentUrl):
|
||||
if document.url.startswith("/media-key/"):
|
||||
# Local file, retrieve from object storage
|
||||
key = document.url[len("/media-key/") :]
|
||||
# Security check: ensure the document belongs to the conversation
|
||||
if not key.startswith(f"{self.conversation.pk}/"):
|
||||
raise ValueError("Document URL does not belong to the conversation.")
|
||||
# Retrieve the document data
|
||||
with default_storage.open(key, "rb") as file:
|
||||
document_data = file.read()
|
||||
parsed_content = document_store.parse_and_store_document(
|
||||
name=document.identifier,
|
||||
content_type=document.media_type,
|
||||
content=document_data,
|
||||
)
|
||||
else:
|
||||
# Remote URL
|
||||
raise ValueError("External document URL are not accepted yet.")
|
||||
else:
|
||||
parsed_content = document_store.parse_and_store_document(
|
||||
name=document.identifier,
|
||||
content_type=document.media_type,
|
||||
content=document.data,
|
||||
)
|
||||
parsed_content = document_store.parse_and_store_document(
|
||||
name=document.identifier,
|
||||
content_type=document.media_type,
|
||||
content=document.data,
|
||||
)
|
||||
await models.ChatConversationContext.objects.acreate(
|
||||
conversation=self.conversation,
|
||||
kind=models.ChatConversationContextKind.DOCUMENT.value,
|
||||
name=document.identifier,
|
||||
content=parsed_content,
|
||||
)
|
||||
|
||||
if not document.media_type.startswith("text/"):
|
||||
md_attachment = await models.ChatConversationAttachment.objects.acreate(
|
||||
conversation=self.conversation,
|
||||
uploaded_by=self.user,
|
||||
key=key or f"{self.conversation.pk}/attachments/{document.identifier}.md",
|
||||
file_name=f"{document.identifier}.md",
|
||||
content_type="text/markdown",
|
||||
conversion_from=key, # might be None
|
||||
)
|
||||
default_storage.save(md_attachment.key, BytesIO(parsed_content.encode("utf8")))
|
||||
md_attachment.upload_state = models.AttachmentStatus.READY
|
||||
await md_attachment.asave(update_fields=["upload_state", "updated_at"])
|
||||
|
||||
def prepare_prompt( # noqa: PLR0912 # pylint: disable=too-many-branches
|
||||
def prepare_prompt(
|
||||
self, message: UIMessage
|
||||
) -> Tuple[str, List[BinaryContent | ImageUrl], List[BinaryContent]]:
|
||||
) -> Tuple[str, List[BinaryContent], List[BinaryContent]]:
|
||||
"""
|
||||
Prepare the user prompt for the agent.
|
||||
|
||||
@@ -329,10 +268,6 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
attachment_images.append(content)
|
||||
else:
|
||||
attachment_documents.append(content)
|
||||
elif isinstance(content, ImageUrl):
|
||||
attachment_images.append(content)
|
||||
elif isinstance(content, DocumentUrl):
|
||||
attachment_documents.append(content)
|
||||
else:
|
||||
# Should never happen, but just in case
|
||||
raise ValueError(f"Unsupported UserContent type: {type(content)}")
|
||||
@@ -357,57 +292,32 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
messages: List[UIMessage],
|
||||
force_web_search: bool = False,
|
||||
) -> events_v4.Event | events_v5.Event:
|
||||
"""
|
||||
Drive the agent for the provided user message, stream Vercel-AI-SDK event parts representing model and tool activity, and persist the final conversation state.
|
||||
|
||||
Parameters:
|
||||
messages (List[UIMessage]): UI messages for the conversation; the last message must be from the user.
|
||||
force_web_search (bool): If true, require the agent to invoke the configured web search tool before answering (ignored if the feature or tool is unavailable).
|
||||
|
||||
Returns:
|
||||
events_v4.Event | events_v5.Event: Streamed event parts such as `TextPart`, `ToolCallPart`/`ToolCallStreamingStartPart`/`ToolCallDeltaPart`, `ToolResultPart`, `ReasoningPart`, `SourcePart`, `DataPart`, `StartStepPart`, and `FinishMessagePart` that drive frontend updates.
|
||||
"""
|
||||
"""Run the Pydantic AI agent and stream events."""
|
||||
if messages[-1].role != "user":
|
||||
return
|
||||
|
||||
# Langfuse settings
|
||||
if self._langfuse_available:
|
||||
if self._store_analytics:
|
||||
langfuse = get_client()
|
||||
langfuse.update_current_trace(
|
||||
session_id=str(self.conversation.pk),
|
||||
user_id=str(self.user.sub),
|
||||
metadata={
|
||||
"user_fqdn": self.user.email.split("@")[-1], # no need for security here
|
||||
},
|
||||
)
|
||||
|
||||
history = ModelMessagesTypeAdapter.validate_python(self.conversation.pydantic_messages)
|
||||
history = update_history_local_urls(
|
||||
self.conversation, history
|
||||
) # presign URLs for local images
|
||||
|
||||
user_prompt, input_images, input_documents = self.prepare_prompt(messages[-1])
|
||||
|
||||
image_key_mapping = {}
|
||||
if input_images:
|
||||
# presign URLs for local images
|
||||
input_images = update_local_urls(
|
||||
self.conversation, input_images, updated_url=image_key_mapping
|
||||
)
|
||||
|
||||
if self._langfuse_available:
|
||||
langfuse.update_current_trace(
|
||||
input=user_prompt if self._store_analytics else "REDACTED"
|
||||
)
|
||||
if self._store_analytics:
|
||||
langfuse.update_current_trace(input=user_prompt)
|
||||
|
||||
usage = {"promptTokens": 0, "completionTokens": 0}
|
||||
|
||||
conversation_has_documents = self._is_document_upload_enabled and (
|
||||
bool(self.conversation.collection_id)
|
||||
or bool(
|
||||
await models.ChatConversationAttachment.objects.filter(
|
||||
await models.ChatConversationContext.objects.filter(
|
||||
conversation=self.conversation,
|
||||
content_type__startswith="text/",
|
||||
kind=models.ChatConversationContextKind.DOCUMENT.value,
|
||||
).aexists()
|
||||
)
|
||||
)
|
||||
@@ -456,28 +366,6 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
|
||||
await self._agent_stop_streaming(force_cache_check=True)
|
||||
|
||||
generated_title = None
|
||||
|
||||
# +1 because we're about to add a new user message
|
||||
current_user_count = sum(1 for msg in self.conversation.messages if msg.role == "user") + 1
|
||||
if (
|
||||
current_user_count == settings.AUTO_TITLE_AFTER_USER_MESSAGES
|
||||
and not self.conversation.title_set_by_user_at
|
||||
):
|
||||
generated_title = await self._generate_title()
|
||||
|
||||
# Notify frontend about the title update
|
||||
if generated_title:
|
||||
yield events_v4.DataPart(
|
||||
data=[
|
||||
{
|
||||
"type": "conversation_metadata",
|
||||
"conversationId": str(self.conversation.pk),
|
||||
"title": generated_title,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
if force_web_search and not self._is_web_search_enabled:
|
||||
logger.warning("Web search is forced but the feature is disabled, ignoring.")
|
||||
force_web_search = False
|
||||
@@ -499,86 +387,22 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
|
||||
_tool_is_streaming = False
|
||||
_model_response_message_id = None
|
||||
|
||||
# Check for existing non-PDF documents in the conversation:
|
||||
# - if no document at all: do nothing
|
||||
# - if only PDFs: prepare document URLs for the agent
|
||||
# - if other document types: add the RAG search tool
|
||||
# to allow searching in all kinds of documents
|
||||
has_not_pdf_docs = await (
|
||||
models.ChatConversationAttachment.objects.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from=""),
|
||||
conversation=self.conversation,
|
||||
)
|
||||
.exclude(
|
||||
Q(content_type__startswith="image/") | Q(content_type="application/pdf"),
|
||||
)
|
||||
.aexists()
|
||||
)
|
||||
|
||||
document_urls = []
|
||||
if not conversation_has_documents and not has_not_pdf_docs:
|
||||
# No documents to process
|
||||
pass
|
||||
elif has_not_pdf_docs:
|
||||
add_document_rag_search_tool(self.conversation_agent)
|
||||
|
||||
@self.conversation_agent.instructions
|
||||
def summarization_system_prompt() -> str:
|
||||
return (
|
||||
"When you receive a result from the summarization tool, you MUST return it "
|
||||
"directly to the user without any modification, paraphrasing, or additional "
|
||||
"summarization."
|
||||
"The tool already produces optimized summaries that should be presented "
|
||||
"verbatim."
|
||||
"You may translate the summary if required, but you MUST preserve all the "
|
||||
"information from the original summary."
|
||||
"You may add a follow-up question after the summary if needed."
|
||||
)
|
||||
|
||||
# Inform the model (system-level) that documents are attached and available
|
||||
@self.conversation_agent.system_prompt
|
||||
def attached_documents_note() -> str:
|
||||
return (
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; consider them already available "
|
||||
"via the internal store."
|
||||
)
|
||||
|
||||
@self.conversation_agent.tool(name="summarize", retries=2)
|
||||
@functools.wraps(document_summarize)
|
||||
async def summarize(ctx: RunContext, *args, **kwargs) -> ToolReturn:
|
||||
"""Wrap the document_summarize tool to provide context and add the tool."""
|
||||
return await document_summarize(ctx, *args, **kwargs)
|
||||
else:
|
||||
conversation_documents = [
|
||||
cd
|
||||
async for cd in models.ChatConversationAttachment.objects.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from=""),
|
||||
conversation=self.conversation,
|
||||
)
|
||||
.exclude(
|
||||
content_type__startswith="image/",
|
||||
)
|
||||
.values_list("key", "content_type")
|
||||
]
|
||||
|
||||
for doc_key, doc_content_type in conversation_documents:
|
||||
if doc_content_type == "application/pdf":
|
||||
_presigned_url = generate_retrieve_policy(doc_key)
|
||||
document_urls.append(
|
||||
DocumentUrl(
|
||||
url=_presigned_url,
|
||||
identifier=doc_key.split("/")[-1],
|
||||
media_type="application/pdf",
|
||||
)
|
||||
)
|
||||
image_key_mapping[_presigned_url] = f"/media-key/{doc_key}"
|
||||
|
||||
async with AsyncExitStack() as stack:
|
||||
# MCP servers (if any) can be initialized here
|
||||
mcp_servers = [await stack.enter_async_context(mcp) for mcp in get_mcp_servers()]
|
||||
|
||||
if conversation_has_documents:
|
||||
add_document_rag_search_tool(self.conversation_agent)
|
||||
|
||||
@self.conversation_agent.system_prompt
|
||||
def summarize_instructions() -> str:
|
||||
"""Dynamic system prompt function to add RAG instructions if any."""
|
||||
return (
|
||||
"If the user wants a summary of document(s), invoke summarize tool "
|
||||
"without asking the user for the document itself. The tool will handle "
|
||||
"any necessary extraction and summarization based on the internal context."
|
||||
)
|
||||
|
||||
_final_output_from_tool = None
|
||||
_ui_sources = []
|
||||
|
||||
@@ -588,10 +412,15 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
history.append(ModelResponse(parts=[TextPart(content="ok")], kind="response"))
|
||||
|
||||
async with self.conversation_agent.iter(
|
||||
[user_prompt] + input_images + document_urls,
|
||||
message_history=history, # history will pass through agent's history_processors
|
||||
[user_prompt] + input_images,
|
||||
message_history=history,
|
||||
deps=self._context_deps,
|
||||
toolsets=mcp_servers,
|
||||
output_type=(
|
||||
[ToolOutput(hand_off_to_summarization_agent, name="summarize"), str]
|
||||
if conversation_has_documents
|
||||
else str
|
||||
),
|
||||
) as run:
|
||||
async for node in run:
|
||||
await self._agent_stop_streaming()
|
||||
@@ -735,13 +564,56 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
logger.error("_model_response_message_id already set")
|
||||
_model_response_message_id = (
|
||||
str(uuid.uuid4())
|
||||
if not self._langfuse_available
|
||||
if not self._store_analytics
|
||||
else f"trace-{langfuse.get_current_trace_id()}"
|
||||
)
|
||||
yield events_v4.StartStepPart(
|
||||
message_id=_model_response_message_id,
|
||||
)
|
||||
|
||||
if (
|
||||
isinstance(node.data, FinalResult)
|
||||
and node.data.tool_name == "summarize"
|
||||
):
|
||||
yield events_v4.ToolResultPart(
|
||||
tool_call_id=node.data.tool_call_id,
|
||||
result={"state": "done"}, # content not needed here
|
||||
)
|
||||
final_output = node.data.output
|
||||
if isinstance(final_output, ToolReturn):
|
||||
_final_output_from_tool = final_output.content
|
||||
yield events_v4.TextPart(text=final_output.content)
|
||||
|
||||
if final_output.metadata and (
|
||||
sources := final_output.metadata.get("sources")
|
||||
):
|
||||
for source_url in sources:
|
||||
url_source = LanguageModelV1Source(
|
||||
sourceType="url",
|
||||
id=str(uuid.uuid4()),
|
||||
url=source_url,
|
||||
providerMetadata={},
|
||||
)
|
||||
_new_source_ui = SourceUIPart(
|
||||
type="source", source=url_source
|
||||
)
|
||||
_ui_sources.append(_new_source_ui)
|
||||
yield events_v4.SourcePart(
|
||||
**_new_source_ui.source.model_dump()
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Unexpected final result type: %s %s",
|
||||
type(final_output),
|
||||
final_output,
|
||||
)
|
||||
|
||||
else:
|
||||
logger.warning(
|
||||
"Unknown node type encountered: %s",
|
||||
type(node),
|
||||
)
|
||||
|
||||
# Final usage summary
|
||||
final_usage = run.usage()
|
||||
usage["promptTokens"] = final_usage.input_tokens
|
||||
@@ -752,18 +624,15 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
# Persist conversation
|
||||
await sync_to_async(self._update_conversation)(
|
||||
final_output=run.result.new_messages(),
|
||||
raw_final_output=run.result.new_messages_json(),
|
||||
usage=usage,
|
||||
final_output_from_tool=_final_output_from_tool,
|
||||
ui_sources=_ui_sources,
|
||||
model_response_message_id=_model_response_message_id,
|
||||
image_key_mapping=image_key_mapping or None,
|
||||
generated_title=generated_title,
|
||||
)
|
||||
|
||||
if self._langfuse_available:
|
||||
langfuse.update_current_trace(
|
||||
output=run.result.output if self._store_analytics else "REDACTED"
|
||||
)
|
||||
if self._store_analytics:
|
||||
langfuse.update_current_trace(output=run.result.output)
|
||||
|
||||
# Vercel finish message
|
||||
yield events_v4.FinishMessagePart(
|
||||
@@ -778,30 +647,24 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
self,
|
||||
*,
|
||||
final_output: List[ModelRequest | ModelMessage],
|
||||
raw_final_output: bytes,
|
||||
usage: Dict[str, int],
|
||||
final_output_from_tool: str | None,
|
||||
ui_sources: List[SourceUIPart] = None,
|
||||
model_response_message_id: str | None = None,
|
||||
image_key_mapping: Dict[str, str] = None,
|
||||
generated_title: str | None = None,
|
||||
): # pylint: disable=too-many-arguments
|
||||
"""
|
||||
Merge the agent's final outputs into the conversation and persist updated conversation state.
|
||||
|
||||
Parameters:
|
||||
final_output (List[ModelRequest | ModelMessage]): Sequence of model requests and responses produced by the agent run; these will be merged into a single request and a single response before saving.
|
||||
usage (Dict[str, int]): Token usage statistics to store on the conversation (e.g., promptTokens, completionTokens).
|
||||
final_output_from_tool (str | None): Optional text produced by a tool that should be appended to the final model response.
|
||||
ui_sources (List[SourceUIPart], optional): Optional UI-visible source parts to attach to the final response message.
|
||||
model_response_message_id (str | None, optional): If provided, assign this id to the saved model response UI message; if omitted, a warning will be logged.
|
||||
image_key_mapping (Dict[str, str], optional): Mapping from original (unsigned) media URLs to presigned/rewritten URLs; applied to image/document references in the merged request parts.
|
||||
generated_title (str | None, optional): Optional auto-generated conversation title to apply to the conversation.
|
||||
|
||||
Behavior:
|
||||
- Merges multiple model request/response objects into a single ModelRequest and ModelResponse.
|
||||
- Rewrites image/document URLs in user prompt parts when an image_key_mapping is provided.
|
||||
- Converts merged model messages to UI messages, appends ui_sources if present, and sets the response message id when supplied.
|
||||
- Appends the merged request and response messages to the conversation, updates agent usage and pydantic messages, applies a generated title if given, and saves the conversation.
|
||||
Save everything related to the conversation.
|
||||
|
||||
Things to improve here:
|
||||
- The way we need to add the UI sources to the final output message.
|
||||
|
||||
Args:
|
||||
final_output (List[ModelRequest | ModelMessage]): The final output from the agent.
|
||||
raw_final_output (bytes): The raw final output in bytes.
|
||||
usage (Dict[str, int]): The token usage statistics.
|
||||
user_initial_prompt_str (str | None): The initial user prompt string, if any.
|
||||
ui_sources (List[SourceUIPart]): Optional UI sources to include in the conversation.
|
||||
"""
|
||||
_merged_final_output_request = ModelRequest(
|
||||
parts=[
|
||||
@@ -817,15 +680,6 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
kind="response",
|
||||
)
|
||||
|
||||
if image_key_mapping:
|
||||
for part in _merged_final_output_request.parts:
|
||||
if isinstance(part, UserPromptPart):
|
||||
for content in part.content:
|
||||
if isinstance(content, (ImageUrl, DocumentUrl)) and (
|
||||
unsigned_url := image_key_mapping.get(content.url)
|
||||
):
|
||||
content.url = unsigned_url
|
||||
|
||||
_output_ui_message = model_message_to_ui_message(_merged_final_output_message)
|
||||
if ui_sources:
|
||||
_output_ui_message.parts += ui_sources
|
||||
@@ -840,50 +694,11 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
]
|
||||
self.conversation.agent_usage = usage
|
||||
|
||||
final_output_json = json.loads(
|
||||
ModelMessagesTypeAdapter.dump_json(final_output).decode("utf-8")
|
||||
logger.debug(
|
||||
"raw_final_output: %s %s",
|
||||
raw_final_output.decode("utf-8"),
|
||||
json.loads(raw_final_output.decode("utf-8")),
|
||||
)
|
||||
logger.debug("final_output_json: %s", final_output_json)
|
||||
self.conversation.pydantic_messages += json.loads(
|
||||
ModelMessagesTypeAdapter.dump_json(final_output).decode("utf-8")
|
||||
)
|
||||
if generated_title:
|
||||
self.conversation.title = generated_title
|
||||
self.conversation.pydantic_messages += json.loads(raw_final_output.decode("utf-8"))
|
||||
|
||||
self.conversation.save()
|
||||
|
||||
async def _generate_title(self) -> str | None:
|
||||
"""
|
||||
Create a concise conversation title based on the conversation's first messages.
|
||||
|
||||
Uses the summarization agent to produce a short title in the same language as the user's messages. Returns the generated title text trimmed to at most 100 characters, or `None` if generation fails or produces no text.
|
||||
|
||||
Returns:
|
||||
str | None: The generated title (trimmed to 100 characters), or `None` when no title is available.
|
||||
"""
|
||||
|
||||
# Build context from the first messages
|
||||
context = "\n".join(
|
||||
f"{msg.role}: {msg.content[:300]}" # Limit content length per message
|
||||
for msg in self.conversation.messages[:6] # First few messages (3 user + 3 assistant)
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"Generate a short, concise title (maximum 60 characters) for this conversation. "
|
||||
"The title should capture the main topic or intent. "
|
||||
"Return ONLY the title text, nothing else. No quotes, no explanations.\n\n"
|
||||
"Return the title text in the same language the user messages are written."
|
||||
f"If in doubt, use {self.language or 'French'}."
|
||||
f"Conversation:\n{context}"
|
||||
)
|
||||
|
||||
try:
|
||||
agent = SummarizationAgent()
|
||||
result = await agent.run(prompt)
|
||||
title = (result.output or "").strip()[:100] # Enforce max length
|
||||
logger.info("Generated title for conversation %s: %s", self.conversation.pk, title)
|
||||
return title if title else None
|
||||
except Exception as exc: # pylint: disable=broad-except #noqa: BLE001
|
||||
logger.warning(
|
||||
"Failed to generate title for conversation %s: %s", self.conversation.pk, exc
|
||||
)
|
||||
return None
|
||||
@@ -12,8 +12,6 @@ from typing import List
|
||||
|
||||
from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
DocumentUrl,
|
||||
ImageUrl,
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
@@ -71,27 +69,15 @@ def ui_message_to_user_content(message: UIMessage) -> List[UserContent]:
|
||||
identifier=experimental_attachment.name,
|
||||
)
|
||||
)
|
||||
elif experimental_attachment.contentType.startswith("image/"):
|
||||
user_contents.append(
|
||||
ImageUrl(
|
||||
url=experimental_attachment.url,
|
||||
media_type=experimental_attachment.contentType,
|
||||
identifier=experimental_attachment.name,
|
||||
)
|
||||
)
|
||||
else:
|
||||
user_contents.append(
|
||||
DocumentUrl(
|
||||
url=experimental_attachment.url,
|
||||
media_type=experimental_attachment.contentType,
|
||||
identifier=experimental_attachment.name,
|
||||
)
|
||||
raise ValueError(
|
||||
f"Unsupported experimental attachment URL format: {experimental_attachment.url}"
|
||||
)
|
||||
|
||||
return user_contents
|
||||
|
||||
|
||||
def model_message_to_ui_message(model_message: ModelMessage) -> UIMessage: # noqa: PLR0912, PLR0915 # pylint: disable=too-many-statements
|
||||
def model_message_to_ui_message(model_message: ModelMessage) -> UIMessage: # noqa: PLR0912
|
||||
"""
|
||||
Convert a ModelMessage (ModelRequest or ModelResponse) to a UIMessage.
|
||||
"""
|
||||
@@ -129,23 +115,7 @@ def model_message_to_ui_message(model_message: ModelMessage) -> UIMessage: # no
|
||||
+ base64.b64encode(c.data).decode("utf-8"),
|
||||
)
|
||||
)
|
||||
elif isinstance(c, ImageUrl):
|
||||
experimental_attachments.append(
|
||||
Attachment(
|
||||
contentType=c.media_type,
|
||||
url=c.url,
|
||||
name=c.identifier,
|
||||
)
|
||||
)
|
||||
elif isinstance(c, DocumentUrl):
|
||||
experimental_attachments.append(
|
||||
Attachment(
|
||||
contentType=c.media_type,
|
||||
url=c.url,
|
||||
name=c.identifier,
|
||||
)
|
||||
)
|
||||
else: # AudioUrl, VideoUrl
|
||||
else: # ImageUrl, AudioUrl, VideoUrl, DocumentUrl, BinaryContent
|
||||
raise ValueError(
|
||||
f"Unsupported UserContent in UserPromptPart: {type(c)}"
|
||||
)
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
"""Factories for chat application."""
|
||||
|
||||
from uuid import uuid4
|
||||
|
||||
import factory.django
|
||||
|
||||
from core.factories import UserFactory
|
||||
@@ -16,18 +14,3 @@ class ChatConversationFactory(factory.django.DjangoModelFactory):
|
||||
|
||||
class Meta:
|
||||
model = models.ChatConversation
|
||||
|
||||
|
||||
class ChatConversationAttachmentFactory(factory.django.DjangoModelFactory):
|
||||
"""Factory for creating ChatConversationAttachment instances."""
|
||||
|
||||
conversation = factory.SubFactory(ChatConversationFactory)
|
||||
uploaded_by = factory.SubFactory(UserFactory)
|
||||
key = factory.LazyAttribute(
|
||||
lambda obj: f"{obj.conversation.pk}/attachments/{uuid4()}.{obj.file_name.split('.')[-1]}"
|
||||
)
|
||||
file_name = factory.Faker("file_name")
|
||||
content_type = factory.Faker("mime_type")
|
||||
|
||||
class Meta:
|
||||
model = models.ChatConversationAttachment
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Malware detection callbacks"""
|
||||
|
||||
import logging
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
|
||||
from chat.models import ChatConversationAttachment
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
security_logger = logging.getLogger("conversations.security")
|
||||
|
||||
|
||||
def conversation_safe_attachment_callback(file_path, *, conversation_id, **kwargs):
|
||||
"""Callback when a malware scan is completed and unsafe for a conversation attachment."""
|
||||
logger.info("File %s for conversation %s is safe", file_path, conversation_id)
|
||||
|
||||
ChatConversationAttachment.objects.filter(
|
||||
conversation_id=conversation_id, key=file_path
|
||||
).update(upload_state=AttachmentStatus.READY)
|
||||
|
||||
|
||||
def unknown_attachment_callback(file_path, error_info, *, conversation_id, **kwargs) -> bool:
|
||||
"""Callback when a malware scan is completed and unknown for a conversation attachment."""
|
||||
security_logger.warning(
|
||||
"File %s for conversation %s has an unknown reportstatus. Error info: %s",
|
||||
file_path,
|
||||
conversation_id,
|
||||
error_info,
|
||||
)
|
||||
|
||||
error_code = error_info.get("error_code")
|
||||
if error_code == 413:
|
||||
ChatConversationAttachment.objects.filter(
|
||||
conversation_id=conversation_id, key=file_path
|
||||
).update(upload_state=AttachmentStatus.FILE_TOO_LARGE_TO_ANALYZE)
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def conversation_unsafe_attachment_callback(file_path, error_info, *, conversation_id, **kwargs):
|
||||
"""Callback when a malware scan is completed and unsafe for a conversation attachment."""
|
||||
security_logger.warning(
|
||||
"File %s for conversation %s is infected with malware. Error info: %s",
|
||||
file_path,
|
||||
conversation_id,
|
||||
error_info,
|
||||
)
|
||||
|
||||
ChatConversationAttachment.objects.filter(
|
||||
conversation_id=conversation_id, key=file_path
|
||||
).update(upload_state=AttachmentStatus.SUSPICIOUS)
|
||||
@@ -1,100 +0,0 @@
|
||||
# Generated by Django 5.2.7 on 2025-10-17 16:10
|
||||
|
||||
import uuid
|
||||
|
||||
import django.db.models.deletion
|
||||
from django.conf import settings
|
||||
from django.db import migrations, models
|
||||
|
||||
import core.file_upload.enums
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
dependencies = [
|
||||
("chat", "0003_chatconversationcontext"),
|
||||
migrations.swappable_dependency(settings.AUTH_USER_MODEL),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.CreateModel(
|
||||
name="ChatConversationAttachment",
|
||||
fields=[
|
||||
(
|
||||
"id",
|
||||
models.UUIDField(
|
||||
default=uuid.uuid4,
|
||||
editable=False,
|
||||
help_text="primary key for the record as UUID",
|
||||
primary_key=True,
|
||||
serialize=False,
|
||||
verbose_name="id",
|
||||
),
|
||||
),
|
||||
(
|
||||
"created_at",
|
||||
models.DateTimeField(
|
||||
auto_now_add=True,
|
||||
help_text="date and time at which a record was created",
|
||||
verbose_name="created on",
|
||||
),
|
||||
),
|
||||
(
|
||||
"updated_at",
|
||||
models.DateTimeField(
|
||||
auto_now=True,
|
||||
help_text="date and time at which a record was last updated",
|
||||
verbose_name="updated on",
|
||||
),
|
||||
),
|
||||
(
|
||||
"upload_state",
|
||||
models.CharField(
|
||||
choices=core.file_upload.enums.AttachmentStatus.choices,
|
||||
default=core.file_upload.enums.AttachmentStatus["PENDING"],
|
||||
max_length=40,
|
||||
),
|
||||
),
|
||||
(
|
||||
"key",
|
||||
models.CharField(help_text="File path of the attachment in the object storage"),
|
||||
),
|
||||
("file_name", models.CharField(help_text="Original name of the attachment file")),
|
||||
(
|
||||
"content_type",
|
||||
models.CharField(help_text="MIME type of the attachment file", max_length=100),
|
||||
),
|
||||
("size", models.PositiveBigIntegerField(blank=True, null=True)),
|
||||
(
|
||||
"conversation",
|
||||
models.ForeignKey(
|
||||
on_delete=django.db.models.deletion.CASCADE,
|
||||
related_name="attachments",
|
||||
to="chat.chatconversation",
|
||||
),
|
||||
),
|
||||
(
|
||||
"uploaded_by",
|
||||
models.ForeignKey(
|
||||
help_text="User who uploaded the attachment",
|
||||
on_delete=django.db.models.deletion.PROTECT,
|
||||
related_name="uploaded_attachments",
|
||||
to=settings.AUTH_USER_MODEL,
|
||||
),
|
||||
),
|
||||
(
|
||||
"conversion_from",
|
||||
models.CharField(
|
||||
blank=True,
|
||||
help_text="Original file key if the Markdown from another file",
|
||||
null=True,
|
||||
),
|
||||
),
|
||||
],
|
||||
options={
|
||||
"abstract": False,
|
||||
},
|
||||
),
|
||||
migrations.DeleteModel(
|
||||
name="ChatConversationContext",
|
||||
),
|
||||
]
|
||||
@@ -1,21 +0,0 @@
|
||||
# Generated by Django 5.2.9 on 2025-12-30 09:44
|
||||
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
dependencies = [
|
||||
("chat", "0004_chatconversationattachment_and_more"),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AddField(
|
||||
model_name="chatconversation",
|
||||
name="title_set_by_user_at",
|
||||
field=models.DateTimeField(
|
||||
blank=True,
|
||||
help_text="Timestamp when the user manually set the title. If set, prevent automatic title generation.",
|
||||
null=True,
|
||||
),
|
||||
),
|
||||
]
|
||||
+27
-43
@@ -7,7 +7,6 @@ from django.db import models
|
||||
|
||||
from django_pydantic_field import SchemaField
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.models import BaseModel
|
||||
|
||||
from chat.ai_sdk_types import UIMessage
|
||||
@@ -44,12 +43,7 @@ class ChatConversation(BaseModel):
|
||||
null=True,
|
||||
help_text="Title of the chat conversation",
|
||||
)
|
||||
title_set_by_user_at = models.DateTimeField(
|
||||
blank=True,
|
||||
null=True,
|
||||
help_text="Timestamp when the user manually set the title. If set, prevent automatic "
|
||||
"title generation.",
|
||||
)
|
||||
|
||||
ui_messages = models.JSONField(
|
||||
default=list,
|
||||
blank=True,
|
||||
@@ -80,59 +74,49 @@ class ChatConversation(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class ChatConversationAttachment(BaseModel):
|
||||
class ChatConversationContextKind(models.TextChoices):
|
||||
"""Enumeration of chat conversation context kinds."""
|
||||
|
||||
IMAGE = "image" # Context related to an image in base64 format
|
||||
DOCUMENT = "document"
|
||||
|
||||
|
||||
class ChatConversationContext(BaseModel):
|
||||
"""
|
||||
Model representing an attachment associated with a chat conversation.
|
||||
|
||||
This model stores the details of an attachment:
|
||||
- `conversation`: The conversation this attachment belongs to.
|
||||
- `uploaded_by`: The user who uploaded the attachment.
|
||||
- `key`: The file path of the attachment in the object storage.
|
||||
- `file_name`: The original name of the attachment file.
|
||||
- `content_type`: The MIME type of the attachment file.
|
||||
Model representing a chat conversation context.
|
||||
|
||||
This model stores the details of a chat conversation context:
|
||||
- `conversation`: The conversation this context belongs to.
|
||||
- `kind`: The kind of context (e.g., 'image', 'document').
|
||||
- `name`: The key of the context.
|
||||
- `content`: The value of the context.
|
||||
"""
|
||||
|
||||
conversation = models.ForeignKey(
|
||||
ChatConversation,
|
||||
related_name="attachments",
|
||||
related_name="contexts",
|
||||
on_delete=models.CASCADE,
|
||||
null=False,
|
||||
blank=False,
|
||||
)
|
||||
uploaded_by = models.ForeignKey(
|
||||
User,
|
||||
related_name="uploaded_attachments",
|
||||
on_delete=models.PROTECT,
|
||||
null=False,
|
||||
blank=False,
|
||||
help_text="User who uploaded the attachment",
|
||||
)
|
||||
upload_state = models.CharField(
|
||||
max_length=40,
|
||||
choices=AttachmentStatus.choices,
|
||||
default=AttachmentStatus.PENDING,
|
||||
)
|
||||
key = models.CharField(
|
||||
kind = models.CharField(
|
||||
max_length=50,
|
||||
blank=False,
|
||||
null=False,
|
||||
help_text="File path of the attachment in the object storage",
|
||||
help_text="Kind of the chat conversation context (e.g., 'image', 'document')",
|
||||
choices=ChatConversationContextKind,
|
||||
)
|
||||
file_name = models.CharField(
|
||||
blank=False,
|
||||
null=False,
|
||||
help_text="Original name of the attachment file",
|
||||
)
|
||||
content_type = models.CharField(
|
||||
name = models.CharField(
|
||||
max_length=100,
|
||||
blank=False,
|
||||
null=False,
|
||||
help_text="MIME type of the attachment file",
|
||||
help_text="Key of the chat conversation context",
|
||||
)
|
||||
size = models.PositiveBigIntegerField(null=True, blank=True)
|
||||
|
||||
conversion_from = models.CharField(
|
||||
content = models.TextField(
|
||||
blank=True,
|
||||
null=True,
|
||||
help_text="Original file key if the Markdown from another file",
|
||||
help_text="Value of the chat conversation context",
|
||||
)
|
||||
|
||||
class Meta:
|
||||
unique_together = ("conversation", "name")
|
||||
|
||||
@@ -1,18 +1,11 @@
|
||||
"""Serializers for chat application."""
|
||||
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from django.conf import settings
|
||||
from django.utils import timezone
|
||||
|
||||
from django_pydantic_field.rest_framework import SchemaField # pylint: disable=no-name-in-module
|
||||
from drf_spectacular.utils import extend_schema_field
|
||||
from rest_framework import serializers
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.utils import generate_upload_policy
|
||||
|
||||
from chat import models
|
||||
from chat.ai_sdk_types import UIMessage
|
||||
|
||||
@@ -28,20 +21,6 @@ class ChatConversationSerializer(serializers.ModelSerializer):
|
||||
fields = ["id", "title", "created_at", "updated_at", "messages", "owner"]
|
||||
read_only_fields = ["id", "created_at", "updated_at", "messages"]
|
||||
|
||||
def update(self, instance, validated_data):
|
||||
# If title is being changed, mark it as user-set
|
||||
"""
|
||||
Update the ChatConversation instance and record when the title is changed by the user.
|
||||
|
||||
If `validated_data` contains a `title` different from the instance's current title, sets `title_set_by_user_at` to the current time.
|
||||
|
||||
Returns:
|
||||
The updated ChatConversation instance.
|
||||
"""
|
||||
if "title" in validated_data and validated_data["title"] != instance.title:
|
||||
instance.title_set_by_user_at = timezone.now()
|
||||
return super().update(instance, validated_data)
|
||||
|
||||
|
||||
class ChatConversationInputSerializer(serializers.Serializer):
|
||||
"""
|
||||
@@ -163,54 +142,3 @@ class LLMConfigurationSerializer(serializers.Serializer): # pylint: disable=abs
|
||||
"""Serializer for LLM configuration."""
|
||||
|
||||
models = LLModelSerializer(many=True)
|
||||
|
||||
|
||||
class ChatConversationAttachmentSerializer(serializers.ModelSerializer):
|
||||
"""Serializer for chat conversation attachments."""
|
||||
|
||||
url = serializers.SerializerMethodField()
|
||||
|
||||
class Meta: # pylint: disable=missing-class-docstring
|
||||
model = models.ChatConversationAttachment
|
||||
fields = ["id", "key", "content_type", "file_name", "size", "upload_state", "url"]
|
||||
read_only_fields = ["id", "key", "content_type", "file_name", "size", "upload_state"]
|
||||
|
||||
def get_url(self, attachment) -> str | None:
|
||||
"""Return the URL of the attachment."""
|
||||
if attachment.upload_state not in (
|
||||
AttachmentStatus.FILE_TOO_LARGE_TO_ANALYZE,
|
||||
AttachmentStatus.SUSPICIOUS,
|
||||
AttachmentStatus.READY,
|
||||
):
|
||||
return None
|
||||
|
||||
return f"{settings.MEDIA_BASE_URL}{settings.MEDIA_URL}{quote(attachment.key)}"
|
||||
|
||||
|
||||
class CreateChatConversationAttachmentSerializer(serializers.ModelSerializer):
|
||||
"""Serializer for creating chat conversation attachments."""
|
||||
|
||||
policy = serializers.SerializerMethodField()
|
||||
uploaded_by = serializers.HiddenField(default=serializers.CurrentUserDefault())
|
||||
key = serializers.CharField(read_only=True) # Key is generated server-side
|
||||
|
||||
class Meta: # pylint: disable=missing-class-docstring
|
||||
model = models.ChatConversationAttachment
|
||||
fields = ["id", "key", "content_type", "file_name", "size", "policy", "uploaded_by"]
|
||||
|
||||
def get_policy(self, attachment) -> str:
|
||||
"""Return the policy to use if the item is a file."""
|
||||
return generate_upload_policy(attachment.key)
|
||||
|
||||
def validate_size(self, size: Optional[int]) -> Optional[int]:
|
||||
"""Validate that the size is not greater than the maximum allowed size."""
|
||||
if not size:
|
||||
return size
|
||||
|
||||
if size > settings.ATTACHMENT_MAX_SIZE:
|
||||
max_size = settings.ATTACHMENT_MAX_SIZE // (1024 * 1024)
|
||||
raise serializers.ValidationError(
|
||||
f"File size exceeds the maximum limit of {max_size:d} MB."
|
||||
)
|
||||
|
||||
return size
|
||||
-66
@@ -1,66 +0,0 @@
|
||||
"""Unit tests for add_document_rag_search_tool_from_setting integration with AIAgentService."""
|
||||
|
||||
import pytest
|
||||
|
||||
from core.factories import UserFactory
|
||||
|
||||
from chat.clients.pydantic_ai import AIAgentService
|
||||
from chat.factories import ChatConversationFactory
|
||||
from chat.llm_configuration import LLModel, LLMProvider
|
||||
|
||||
pytestmark = pytest.mark.django_db()
|
||||
|
||||
|
||||
def test_ai_agent_service_adds_rag_tools_from_settings(settings):
|
||||
"""Test that AIAgentService adds RAG tools from SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS."""
|
||||
settings.LLM_CONFIGURATIONS = {
|
||||
"default-model": LLModel(
|
||||
hrid="default-model",
|
||||
model_name="amazing-llm",
|
||||
human_readable_name="Amazing LLM",
|
||||
is_active=True,
|
||||
icon=None,
|
||||
system_prompt="You are an amazing assistant.",
|
||||
tools=[],
|
||||
provider=LLMProvider(hrid="unused", base_url="https://example.com", api_key="key"),
|
||||
),
|
||||
}
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
"french_public_services": {
|
||||
"collection_ids": [784, 785],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": (
|
||||
"Use this tool when the user asks for information about French public services."
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
user = UserFactory()
|
||||
conversation = ChatConversationFactory(owner=user)
|
||||
|
||||
# Create the service
|
||||
service = AIAgentService(conversation, user=user)
|
||||
|
||||
# Check that tools were added to the conversation_agent
|
||||
agent_tools = service.conversation_agent._function_toolset.tools # pylint: disable=protected-access
|
||||
|
||||
assert "legal_documents" in agent_tools
|
||||
assert "french_public_services" in agent_tools
|
||||
|
||||
# Verify tool names and descriptions
|
||||
assert agent_tools["legal_documents"].name == "legal_documents"
|
||||
assert (
|
||||
agent_tools["legal_documents"].description
|
||||
== "Use this tool to search legal documents and laws."
|
||||
)
|
||||
|
||||
assert agent_tools["french_public_services"].name == "french_public_services"
|
||||
assert (
|
||||
agent_tools["french_public_services"].description
|
||||
== "Use this tool when the user asks for information about French public services."
|
||||
)
|
||||
@@ -1,270 +0,0 @@
|
||||
"""Unit tests for Langfuse tracing in AIAgentService."""
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
from asgiref.sync import sync_to_async
|
||||
from langfuse import Langfuse
|
||||
from pydantic_ai.messages import ModelMessage
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
|
||||
from core.factories import UserFactory
|
||||
|
||||
from chat.ai_sdk_types import TextUIPart, UIMessage
|
||||
from chat.clients.pydantic_ai import AIAgentService
|
||||
from chat.factories import ChatConversationFactory
|
||||
|
||||
pytestmark = pytest.mark.django_db()
|
||||
|
||||
|
||||
@pytest.fixture(name="langfuse_client", scope="function")
|
||||
def langfuse_client_fixture():
|
||||
"""Fixture to init langfuse for tests."""
|
||||
langfuse_client = Langfuse(
|
||||
public_key="pk-test-key",
|
||||
secret_key="sk-test-key",
|
||||
host="https://langfuse.example.com",
|
||||
environment="test",
|
||||
debug=True,
|
||||
)
|
||||
yield langfuse_client
|
||||
langfuse_client._resources.prompt_cache._task_manager.shutdown() # pylint: disable=protected-access
|
||||
langfuse_client.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def base_settings(settings):
|
||||
"""Set up base settings for the tests."""
|
||||
settings.AI_BASE_URL = "https://api.llm.com/v1/"
|
||||
settings.AI_API_KEY = "test-key"
|
||||
settings.AI_MODEL = "model-123"
|
||||
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful assistant"
|
||||
settings.AI_AGENT_TOOLS = []
|
||||
|
||||
|
||||
@pytest.fixture(name="ui_messages")
|
||||
def ui_messages_fixture():
|
||||
"""Fixture for test UI messages."""
|
||||
return [
|
||||
UIMessage(
|
||||
id="msg-1",
|
||||
role="user",
|
||||
content="Hello, how are you?",
|
||||
parts=[TextUIPart(type="text", text="Hello, how are you?")],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(name="agent_model")
|
||||
def agent_model_fixture():
|
||||
"""Fixture for agent model function."""
|
||||
|
||||
async def _agent_model(_messages: list[ModelMessage], _info: AgentInfo):
|
||||
"""Simple agent model that returns a fixed response."""
|
||||
yield "Hello! I'm doing well, thank you for asking."
|
||||
|
||||
return FunctionModel(stream_function=_agent_model)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@responses.activate
|
||||
async def test_langfuse_span_created_when_enabled_and_analytics_allowed(
|
||||
agent_model, ui_messages, settings, langfuse_client
|
||||
):
|
||||
"""Test Langfuse span is created when enabled and user allows analytics."""
|
||||
settings.LANGFUSE_ENABLED = True
|
||||
|
||||
# Mock Langfuse HTTP endpoints
|
||||
responses.add(
|
||||
responses.POST,
|
||||
"https://langfuse.example.com/api/public/otel/v1/traces",
|
||||
json={"success": True},
|
||||
status=200,
|
||||
)
|
||||
|
||||
# Create user with analytics enabled
|
||||
user = await sync_to_async(UserFactory)(allow_conversation_analytics=True)
|
||||
conversation = await sync_to_async(ChatConversationFactory)(owner=user)
|
||||
|
||||
service = AIAgentService(conversation, user=user)
|
||||
results = []
|
||||
with service.conversation_agent.override(model=agent_model):
|
||||
async for result in service.stream_text_async(ui_messages):
|
||||
results.append(result)
|
||||
|
||||
# Verify that results were generated
|
||||
assert results == ["Hello! I'm doing well, thank you for asking."]
|
||||
|
||||
langfuse_client.flush()
|
||||
|
||||
# Verify Langfuse HTTP calls were made
|
||||
assert len(responses.calls) == 1
|
||||
assert (
|
||||
responses.calls[0].request.url == "https://langfuse.example.com/api/public/otel/v1/traces"
|
||||
)
|
||||
|
||||
# quite complex to parse the full body, so just check that expected output is in there
|
||||
assert b"Hello! I'm doing well, thank you for asking." in responses.calls[0].request.body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@responses.activate
|
||||
async def test_langfuse_span_created_when_enabled_and_analytics_disabled(
|
||||
agent_model, ui_messages, settings, langfuse_client
|
||||
):
|
||||
"""Test Langfuse span is created even when user disallows analytics."""
|
||||
settings.LANGFUSE_ENABLED = True
|
||||
|
||||
# Mock Langfuse HTTP endpoints
|
||||
responses.add(
|
||||
responses.POST,
|
||||
"https://langfuse.example.com/api/public/otel/v1/traces",
|
||||
json={"success": True},
|
||||
status=200,
|
||||
)
|
||||
|
||||
# Create user with analytics disabled
|
||||
user = await sync_to_async(UserFactory)(allow_conversation_analytics=False)
|
||||
conversation = await sync_to_async(ChatConversationFactory)(owner=user)
|
||||
|
||||
service = AIAgentService(conversation, user=user)
|
||||
results = []
|
||||
with service.conversation_agent.override(model=agent_model):
|
||||
async for result in service.stream_text_async(ui_messages):
|
||||
results.append(result)
|
||||
|
||||
# Verify that results were generated
|
||||
assert results == ["Hello! I'm doing well, thank you for asking."]
|
||||
|
||||
langfuse_client.flush()
|
||||
|
||||
# Verify Langfuse HTTP calls were made
|
||||
assert len(responses.calls) == 1
|
||||
assert (
|
||||
responses.calls[0].request.url == "https://langfuse.example.com/api/public/otel/v1/traces"
|
||||
)
|
||||
|
||||
# quite complex to parse the full body, so just check that expected output is in there
|
||||
assert b"Hello! I'm doing well, thank you for asking." not in responses.calls[0].request.body
|
||||
assert b"REDACTED" in responses.calls[0].request.body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@responses.activate
|
||||
async def test_no_langfuse_span_when_disabled(agent_model, ui_messages, settings, langfuse_client):
|
||||
"""Test Langfuse span is not created when Langfuse is disabled."""
|
||||
settings.LANGFUSE_ENABLED = False
|
||||
|
||||
# Mock Langfuse HTTP endpoints (should not be called)
|
||||
responses.add(
|
||||
responses.POST,
|
||||
"https://langfuse.example.com/api/public/ingestion",
|
||||
json={"success": True},
|
||||
status=200,
|
||||
)
|
||||
|
||||
user = await sync_to_async(UserFactory)(allow_conversation_analytics=True)
|
||||
conversation = await sync_to_async(ChatConversationFactory)(owner=user)
|
||||
|
||||
service = AIAgentService(conversation, user=user)
|
||||
results = []
|
||||
with service.conversation_agent.override(model=agent_model):
|
||||
async for result in service.stream_text_async(ui_messages):
|
||||
results.append(result)
|
||||
|
||||
# Verify that results were generated
|
||||
assert results == ["Hello! I'm doing well, thank you for asking."]
|
||||
|
||||
langfuse_client.flush()
|
||||
|
||||
# Verify NO Langfuse HTTP calls were made
|
||||
assert len(responses.calls) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instrumentation_settings_with_analytics_enabled(settings):
|
||||
"""Test service correctly sets flags when Langfuse and analytics are enabled."""
|
||||
# pylint: disable=protected-access
|
||||
settings.LANGFUSE_ENABLED = True
|
||||
|
||||
user = await sync_to_async(UserFactory)(allow_conversation_analytics=True)
|
||||
conversation = await sync_to_async(ChatConversationFactory)(owner=user)
|
||||
service = AIAgentService(conversation, user=user)
|
||||
|
||||
# Verify that flags are set correctly
|
||||
assert service._langfuse_available is True
|
||||
assert service._store_analytics is True
|
||||
# ConversationAgent should be created successfully
|
||||
assert service.conversation_agent is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instrumentation_settings_with_analytics_disabled(settings):
|
||||
"""Test service correctly sets flags when Langfuse enabled but analytics disabled."""
|
||||
# pylint: disable=protected-access
|
||||
settings.LANGFUSE_ENABLED = True
|
||||
|
||||
user = await sync_to_async(UserFactory)(allow_conversation_analytics=False)
|
||||
conversation = await sync_to_async(ChatConversationFactory)(owner=user)
|
||||
service = AIAgentService(conversation, user=user)
|
||||
|
||||
# Verify that flags are set correctly
|
||||
assert service._langfuse_available is True
|
||||
assert service._store_analytics is False
|
||||
# ConversationAgent should be created successfully
|
||||
assert service.conversation_agent is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_instrumentation_disabled_when_langfuse_disabled(settings):
|
||||
"""Test service correctly sets flags when Langfuse is disabled."""
|
||||
# pylint: disable=protected-access
|
||||
settings.LANGFUSE_ENABLED = False
|
||||
|
||||
user = await sync_to_async(UserFactory)(allow_conversation_analytics=True)
|
||||
conversation = await sync_to_async(ChatConversationFactory)(owner=user)
|
||||
service = AIAgentService(conversation, user=user)
|
||||
|
||||
# Verify that flags are set correctly
|
||||
assert service._langfuse_available is False
|
||||
assert service._store_analytics is False
|
||||
# ConversationAgent should be created successfully
|
||||
assert service.conversation_agent is not None
|
||||
|
||||
|
||||
def test_store_analytics_flag_when_langfuse_enabled_and_user_allows(settings):
|
||||
"""Test _store_analytics is True when Langfuse enabled and user allows analytics."""
|
||||
# pylint: disable=protected-access
|
||||
settings.LANGFUSE_ENABLED = True
|
||||
|
||||
user = UserFactory(allow_conversation_analytics=True)
|
||||
conversation = ChatConversationFactory(owner=user)
|
||||
|
||||
service = AIAgentService(conversation, user=user)
|
||||
assert service._langfuse_available is True
|
||||
assert service._store_analytics is True
|
||||
|
||||
|
||||
def test_store_analytics_flag_when_langfuse_enabled_and_user_disallows(settings):
|
||||
"""Test _store_analytics is False when Langfuse enabled but user disallows analytics."""
|
||||
# pylint: disable=protected-access
|
||||
settings.LANGFUSE_ENABLED = True
|
||||
|
||||
user = UserFactory(allow_conversation_analytics=False)
|
||||
conversation = ChatConversationFactory(owner=user)
|
||||
|
||||
service = AIAgentService(conversation, user=user)
|
||||
assert service._langfuse_available is True
|
||||
assert service._store_analytics is False
|
||||
|
||||
|
||||
def test_store_analytics_flag_when_langfuse_disabled(settings):
|
||||
"""Test _store_analytics is False when Langfuse is disabled."""
|
||||
# pylint: disable=protected-access
|
||||
settings.LANGFUSE_ENABLED = False
|
||||
|
||||
user = UserFactory(allow_conversation_analytics=True)
|
||||
conversation = ChatConversationFactory(owner=user)
|
||||
|
||||
service = AIAgentService(conversation, user=user)
|
||||
assert service._langfuse_available is False
|
||||
assert service._store_analytics is False
|
||||
+17
-72
@@ -3,13 +3,12 @@
|
||||
import datetime
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from pydantic_ai import ImageUrl
|
||||
from pydantic_ai.messages import (
|
||||
AudioUrl,
|
||||
BinaryContent,
|
||||
@@ -37,22 +36,27 @@ from chat.ai_sdk_types import (
|
||||
from chat.clients.pydantic_ui_message_converter import model_message_to_ui_message
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_uuid4_fixture():
|
||||
"""Fixture to mock UUID generation for testing."""
|
||||
with patch("uuid.uuid4", return_value=uuid.UUID("f0cc3bb5-f207-401b-8281-4cba6202991d")):
|
||||
yield
|
||||
|
||||
|
||||
def test_model_message_to_ui_message_text_user_full():
|
||||
"""Test converting a ModelRequest with UserPromptPart containing text to UIMessage."""
|
||||
timestamp = datetime.datetime.now()
|
||||
model_message = ModelRequest(
|
||||
parts=[UserPromptPart(content="Hello!", timestamp=timestamp)], kind="request"
|
||||
)
|
||||
result = model_message_to_ui_message(model_message)
|
||||
|
||||
expected = UIMessage(
|
||||
id=result.id, # Use the generated ID
|
||||
id="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
role="user",
|
||||
content="Hello!",
|
||||
parts=[TextUIPart(type="text", text="Hello!")],
|
||||
createdAt=timestamp,
|
||||
)
|
||||
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@@ -60,15 +64,14 @@ def test_model_message_to_ui_message_text_user_full():
|
||||
def test_model_message_to_ui_message_text_assistant_full():
|
||||
"""Test converting a ModelResponse with TextPart to UIMessage."""
|
||||
model_message = ModelResponse(parts=[TextPart(content="Hi there!")])
|
||||
result = model_message_to_ui_message(model_message)
|
||||
|
||||
expected = UIMessage(
|
||||
id=result.id, # Use the generated ID
|
||||
id="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
role="assistant",
|
||||
content="Hi there!",
|
||||
parts=[TextUIPart(type="text", text="Hi there!")],
|
||||
createdAt=timezone.now(),
|
||||
)
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@@ -79,10 +82,8 @@ def test_model_message_to_ui_message_tool_call_full():
|
||||
model_message = ModelResponse(
|
||||
parts=[ToolCallPart(tool_call_id="id1", tool_name="tool", args=args)]
|
||||
)
|
||||
result = model_message_to_ui_message(model_message)
|
||||
|
||||
expected = UIMessage(
|
||||
id=result.id, # Use the generated ID
|
||||
id="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
role="assistant",
|
||||
content="",
|
||||
parts=[
|
||||
@@ -98,7 +99,7 @@ def test_model_message_to_ui_message_tool_call_full():
|
||||
],
|
||||
createdAt=timezone.now(),
|
||||
)
|
||||
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@@ -107,7 +108,7 @@ def test_model_message_to_ui_message_reasoning_full():
|
||||
"""Test converting a ModelResponse with ThinkingPart to UIMessage."""
|
||||
model_message = ModelResponse(parts=[ThinkingPart(content="reason", signature="sig")])
|
||||
expected = UIMessage(
|
||||
id=str(uuid.uuid4()), # not used in comparison
|
||||
id="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
role="assistant",
|
||||
content="",
|
||||
parts=[
|
||||
@@ -120,7 +121,7 @@ def test_model_message_to_ui_message_reasoning_full():
|
||||
createdAt=timezone.now(),
|
||||
)
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result.id == IsUUID(4)
|
||||
assert result.id == expected.id
|
||||
assert result.role == expected.role
|
||||
assert result.content == expected.content
|
||||
assert result.createdAt == expected.createdAt
|
||||
@@ -161,65 +162,9 @@ def test_model_message_to_ui_message_binary_content():
|
||||
]
|
||||
|
||||
|
||||
def test_model_message_to_ui_message_image_url():
|
||||
"""Test converting a ModelRequest with ImageUrl to UIMessage."""
|
||||
model_message = ModelRequest(
|
||||
parts=[
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"What do you see?",
|
||||
ImageUrl(identifier="doc1.png", url="/media/documents/doc1.png"),
|
||||
]
|
||||
),
|
||||
],
|
||||
kind="request",
|
||||
)
|
||||
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result.role == "user"
|
||||
assert result.parts == [TextUIPart(type="text", text="What do you see?")]
|
||||
assert result.experimental_attachments == [
|
||||
Attachment(
|
||||
name="doc1.png",
|
||||
contentType="image/png",
|
||||
url="/media/documents/doc1.png",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_model_message_to_ui_message_document_url():
|
||||
"""Test converting a ModelRequest with DocumentUrl to UIMessage."""
|
||||
model_message = ModelRequest(
|
||||
parts=[
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"Summarize this",
|
||||
DocumentUrl(
|
||||
identifier="doc1.pdf",
|
||||
url="/media/documents/doc1.pdf",
|
||||
media_type="application/pdf",
|
||||
),
|
||||
]
|
||||
),
|
||||
],
|
||||
kind="request",
|
||||
)
|
||||
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result.role == "user"
|
||||
assert result.parts == [TextUIPart(type="text", text="Summarize this")]
|
||||
assert result.experimental_attachments == [
|
||||
Attachment(
|
||||
name="doc1.pdf",
|
||||
contentType="application/pdf",
|
||||
url="/media/documents/doc1.pdf",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_model_message_to_ui_message_file_parts_full():
|
||||
"""Test handling unsupported file parts in UserPromptPart content."""
|
||||
for part_type in [AudioUrl, VideoUrl]:
|
||||
for part_type in [AudioUrl, VideoUrl, DocumentUrl]:
|
||||
model_message = ModelRequest(
|
||||
parts=[
|
||||
UserPromptPart(
|
||||
|
||||
+8
-8
@@ -3,7 +3,8 @@
|
||||
import base64
|
||||
import datetime
|
||||
|
||||
from pydantic_ai.messages import BinaryContent, DocumentUrl
|
||||
import pytest
|
||||
from pydantic_ai.messages import BinaryContent
|
||||
|
||||
from chat.ai_sdk_types import (
|
||||
Attachment,
|
||||
@@ -208,8 +209,8 @@ def test_user_message_with_reasoning():
|
||||
assert result[0] == "Let me think"
|
||||
|
||||
|
||||
def test_experimental_attachment_url():
|
||||
"""Test conversion of a user message with an experimental attachment URL."""
|
||||
def test_unsupported_experimental_attachment_url():
|
||||
"""Test error handling for unsupported experimental attachment URL format."""
|
||||
ui_message = UIMessage(
|
||||
id="msg8",
|
||||
role="user",
|
||||
@@ -223,11 +224,10 @@ def test_experimental_attachment_url():
|
||||
],
|
||||
)
|
||||
|
||||
result = ui_message_to_user_content(ui_message)
|
||||
assert result == [
|
||||
"Check this file",
|
||||
DocumentUrl(url="https://example.com/file.txt", _media_type="text/plain"),
|
||||
]
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
ui_message_to_user_content(ui_message)
|
||||
|
||||
assert "Unsupported experimental attachment URL format" in str(excinfo.value)
|
||||
|
||||
|
||||
def test_empty_message():
|
||||
|
||||
@@ -4,23 +4,13 @@ import logging
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import formats, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from chat.agents.summarize import SummarizationAgent
|
||||
from chat.clients.pydantic_ai import AIAgentService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@pytest.fixture(name="today_promt_date")
|
||||
def today_prompt_date_fixture():
|
||||
"""Fixture to mock date the system prompt when useless to test it."""
|
||||
_formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
return f"Today is {_formatted_date}."
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_ai_agent_service")
|
||||
def mock_ai_agent_service_fixture():
|
||||
"""Fixture to mock AIAgentService with a custom model."""
|
||||
@@ -49,44 +39,3 @@ def mock_ai_agent_service_fixture():
|
||||
yield
|
||||
|
||||
yield _mock_service
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_summarization_agent")
|
||||
def mock_summarization_agent_fixture():
|
||||
"""Fixture to mock SummarizationAgent with a custom model."""
|
||||
|
||||
@contextmanager
|
||||
def _mock_agent(model):
|
||||
"""Context manager to mock SummarizationAgent with a custom model."""
|
||||
with ExitStack() as stack:
|
||||
|
||||
class SummarizationAgentMock(SummarizationAgent):
|
||||
"""Mocked SummarizationAgent to override the model."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
# We cannot use stack.enter_context(agent.override(model=model))
|
||||
# Because the agent is used outside of this context manager.
|
||||
# So we directly override the protected member.
|
||||
logger.info("Overriding SummarizationAgent model with %s", model)
|
||||
self._model = model # pylint: disable=protected-access
|
||||
|
||||
# Mock the SummarizationAgent in all relevant modules, because first import wins
|
||||
stack.enter_context(
|
||||
patch("chat.agents.summarize.SummarizationAgent", new=SummarizationAgentMock)
|
||||
)
|
||||
stack.enter_context(
|
||||
patch(
|
||||
"chat.tools.document_summarize.SummarizationAgent", new=SummarizationAgentMock
|
||||
)
|
||||
)
|
||||
yield
|
||||
|
||||
yield _mock_agent
|
||||
|
||||
|
||||
PIXEL_PNG = (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00"
|
||||
b"\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\xf8\xff\xff?\x00\x05\xfe\x02\xfe"
|
||||
b"\xa7V\xbd\xfa\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
|
||||
@@ -1,123 +0,0 @@
|
||||
"""Tests for local_media_url_processors."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import (
|
||||
DocumentUrl,
|
||||
ImageUrl,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
TextPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
|
||||
from chat.agents.local_media_url_processors import (
|
||||
update_history_local_urls,
|
||||
update_local_urls,
|
||||
)
|
||||
from chat.factories import ChatConversationFactory
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls(mock_generate_retrieve_policy):
|
||||
"""Test that update_local_urls replaces local URLs with presigned URLs."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_generate_retrieve_policy.side_effect = lambda _key: f"presigned_url-{_key}"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [
|
||||
ImageUrl(url=f"/media-key/{key}"),
|
||||
DocumentUrl(url=f"/media-key/{conversation.pk}/test.pdf"),
|
||||
ImageUrl(url="https://example.com/image.jpg"),
|
||||
]
|
||||
updated_urls = {}
|
||||
|
||||
result = list(update_local_urls(conversation, contents, updated_urls))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].url == f"presigned_url-{key}"
|
||||
assert result[1].url == f"presigned_url-{conversation.pk}/test.pdf"
|
||||
assert result[2].url == "https://example.com/image.jpg"
|
||||
|
||||
assert mock_generate_retrieve_policy.call_count == 2
|
||||
mock_generate_retrieve_policy.assert_any_call(key)
|
||||
mock_generate_retrieve_policy.assert_any_call(f"{conversation.pk}/test.pdf")
|
||||
|
||||
assert len(updated_urls) == 2
|
||||
assert updated_urls[f"presigned_url-{key}"] == f"/media-key/{key}"
|
||||
assert (
|
||||
updated_urls[f"presigned_url-{conversation.pk}/test.pdf"]
|
||||
== f"/media-key/{conversation.pk}/test.pdf"
|
||||
)
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_security_check(mock_generate_retrieve_policy):
|
||||
"""Test that update_local_urls performs a security check."""
|
||||
conversation = ChatConversationFactory()
|
||||
contents = [ImageUrl(url="/media-key/other_conversation/test.jpg")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].url == "/media-key/other_conversation/test.jpg"
|
||||
mock_generate_retrieve_policy.assert_not_called()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_no_media(mock_generate_retrieve_policy):
|
||||
"""Test that update_local_urls handles no media URLs."""
|
||||
conversation = ChatConversationFactory()
|
||||
contents = []
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
assert len(result) == 0
|
||||
mock_generate_retrieve_policy.assert_not_called()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.update_local_urls")
|
||||
def test_update_history_local_urls(mock_update_local_urls):
|
||||
"""Test that update_history_local_urls processes messages."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_update_local_urls.return_value = iter([])
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
user_prompt_content = [
|
||||
ImageUrl(url=f"/media-key/{key}"),
|
||||
DocumentUrl(url=f"/media-key/{conversation.pk}/test.pdf"),
|
||||
]
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart(content=user_prompt_content)]),
|
||||
ModelResponse(parts=[TextPart(content="I see your images.")]),
|
||||
]
|
||||
|
||||
result = update_history_local_urls(conversation, messages)
|
||||
|
||||
assert len(result) == 2
|
||||
mock_update_local_urls.assert_called_once_with(conversation, user_prompt_content)
|
||||
|
||||
|
||||
def test_update_history_local_urls_no_requests():
|
||||
"""Test that update_history_local_urls handles no ModelRequest messages."""
|
||||
conversation = ChatConversationFactory()
|
||||
messages = [
|
||||
ModelRequest(parts=["Hello"]),
|
||||
ModelResponse(parts=[TextPart(content="Hi there!")]),
|
||||
]
|
||||
|
||||
with patch("chat.agents.local_media_url_processors.update_local_urls") as mock:
|
||||
result = update_history_local_urls(conversation, messages)
|
||||
assert result == messages
|
||||
mock.assert_not_called()
|
||||
|
||||
|
||||
def test_update_history_local_urls_no_user_prompt_parts():
|
||||
"""Test that update_history_local_urls handles no UserPromptPart."""
|
||||
conversation = ChatConversationFactory()
|
||||
messages = [ModelRequest(parts=[])]
|
||||
|
||||
with patch("chat.agents.local_media_url_processors.update_local_urls") as mock:
|
||||
result = update_history_local_urls(conversation, messages)
|
||||
assert result == messages
|
||||
mock.assert_not_called()
|
||||
@@ -1,112 +0,0 @@
|
||||
"""Test malware detection callback."""
|
||||
|
||||
from django.core.files.base import ContentFile
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
import pytest
|
||||
from lasuite.malware_detection.enums import ReportStatus
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.malware_detection import malware_detection_callback
|
||||
|
||||
from chat.factories import ChatConversationAttachmentFactory
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
@pytest.fixture(name="safe_file")
|
||||
def fixture_safe_file():
|
||||
"""Create a safe file."""
|
||||
file_path = "test.txt"
|
||||
default_storage.save(file_path, ContentFile("test"))
|
||||
yield file_path
|
||||
default_storage.delete(file_path)
|
||||
|
||||
|
||||
@pytest.fixture(name="unsafe_file")
|
||||
def fixture_unsafe_file():
|
||||
"""Create an unsafe file."""
|
||||
file_path = "unsafe.txt"
|
||||
default_storage.save(file_path, ContentFile("test"))
|
||||
yield file_path
|
||||
default_storage.delete(file_path)
|
||||
|
||||
|
||||
def test_malware_detection_callback_safe_status(safe_file):
|
||||
"""Test malware detection callback with safe status."""
|
||||
|
||||
attachment = ChatConversationAttachmentFactory(key=safe_file)
|
||||
|
||||
malware_detection_callback(
|
||||
safe_file,
|
||||
ReportStatus.SAFE,
|
||||
error_info={},
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=attachment.conversation.pk,
|
||||
)
|
||||
|
||||
assert default_storage.exists(safe_file)
|
||||
attachment.refresh_from_db()
|
||||
assert attachment.upload_state == AttachmentStatus.READY
|
||||
|
||||
|
||||
def test_malware_detection_callback_unsafe_status(unsafe_file):
|
||||
"""Test malware detection callback with unsafe status."""
|
||||
|
||||
attachment = ChatConversationAttachmentFactory(key=unsafe_file)
|
||||
|
||||
malware_detection_callback(
|
||||
unsafe_file,
|
||||
ReportStatus.UNSAFE,
|
||||
error_info={"error": "test", "error_code": 4001},
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=attachment.conversation.pk,
|
||||
)
|
||||
|
||||
assert default_storage.exists(unsafe_file)
|
||||
attachment.refresh_from_db()
|
||||
assert attachment.upload_state == AttachmentStatus.SUSPICIOUS
|
||||
|
||||
|
||||
def test_malware_detection_callback_unknown_status(unsafe_file):
|
||||
"""Test malware detection callback with unknown status."""
|
||||
|
||||
attachment = ChatConversationAttachmentFactory(key=unsafe_file)
|
||||
|
||||
malware_detection_callback(
|
||||
unsafe_file,
|
||||
ReportStatus.UNKNOWN,
|
||||
error_info={"error": "test", "error_code": 4001},
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=attachment.conversation.pk,
|
||||
)
|
||||
|
||||
assert default_storage.exists(unsafe_file)
|
||||
attachment.refresh_from_db()
|
||||
assert attachment.upload_state == AttachmentStatus.SUSPICIOUS
|
||||
|
||||
|
||||
def test_malware_detection_callback_unknown_status_too_large(unsafe_file):
|
||||
"""Test malware detection callback with unknown status and error 413."""
|
||||
|
||||
attachment = ChatConversationAttachmentFactory(key=unsafe_file)
|
||||
|
||||
malware_detection_callback(
|
||||
unsafe_file,
|
||||
ReportStatus.UNKNOWN,
|
||||
error_info={"error": "toooooo big", "error_code": 413},
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=attachment.conversation.pk,
|
||||
)
|
||||
|
||||
assert default_storage.exists(unsafe_file)
|
||||
attachment.refresh_from_db()
|
||||
assert attachment.upload_state == AttachmentStatus.FILE_TOO_LARGE_TO_ANALYZE
|
||||
@@ -1,403 +0,0 @@
|
||||
"""
|
||||
Unit tests for document generic search RAG tool functionality.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import responses
|
||||
import respx
|
||||
from asgiref.sync import sync_to_async
|
||||
from pydantic_ai import Agent, RunContext, RunUsage
|
||||
|
||||
from core.factories import UserFactory
|
||||
|
||||
from chat.tools.document_generic_search_rag import (
|
||||
add_document_rag_search_tool_from_setting,
|
||||
get_specific_rag_search_tool_config,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.django_db()
|
||||
|
||||
|
||||
def test_get_specific_rag_search_tool_config_with_disabled_features(settings):
|
||||
"""Test get_specific_rag_search_tool_config returns tools for enabled features."""
|
||||
user = UserFactory()
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"french_public_services": {
|
||||
"collection_ids": [784, 785],
|
||||
"feature_flag_value": "disabled",
|
||||
"tool_description": (
|
||||
"Use this tool when the user asks for information about French public services, "
|
||||
"the French labor market, employment laws, social benefits, or "
|
||||
"assistance with administrative procedures."
|
||||
),
|
||||
},
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "disabled",
|
||||
"tool_description": "Use this tool to search French legal documents and laws.",
|
||||
"rag_backend_name": "chat.tests.tools.test_document_generic_search_rag.MockRagBackend",
|
||||
},
|
||||
}
|
||||
|
||||
# The fixture tools are disabled by default
|
||||
assert get_specific_rag_search_tool_config(user) == {}
|
||||
|
||||
|
||||
def test_get_specific_rag_search_tool_config_with_enabled_features(settings):
|
||||
"""Test get_specific_rag_search_tool_config returns tools for enabled features."""
|
||||
user = UserFactory()
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"french_public_services": {
|
||||
"collection_ids": [784, 785],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": (
|
||||
"Use this tool when the user asks for information about French public services, "
|
||||
"the French labor market, employment laws, social benefits, or "
|
||||
"assistance with administrative procedures."
|
||||
),
|
||||
},
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search French legal documents and laws.",
|
||||
},
|
||||
}
|
||||
|
||||
assert get_specific_rag_search_tool_config(user) == {
|
||||
"french_public_services": {
|
||||
"collection_ids": [784, 785],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool when the user "
|
||||
"asks for information about "
|
||||
"French public services, the "
|
||||
"French labor market, "
|
||||
"employment laws, social "
|
||||
"benefits, or assistance with "
|
||||
"administrative procedures.",
|
||||
},
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search French legal documents and laws.",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_get_specific_rag_search_tool_config_with_dynamic_features(settings, posthog):
|
||||
"""Test get_specific_rag_search_tool_config with dynamic features."""
|
||||
user = UserFactory()
|
||||
|
||||
responses.post(
|
||||
f"{posthog.host}/flags/?v=2",
|
||||
json={"flags": {"legal-documents": {"enabled": True}}},
|
||||
status=200,
|
||||
)
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"french_public_services": {
|
||||
"collection_ids": [784, 785],
|
||||
"feature_flag_value": "dynamic",
|
||||
"tool_description": (
|
||||
"Use this tool when the user asks for information about French public services, "
|
||||
"the French labor market, employment laws, social benefits, or "
|
||||
"assistance with administrative procedures."
|
||||
),
|
||||
},
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "dynamic",
|
||||
"tool_description": "Use this tool to search French legal documents and laws.",
|
||||
},
|
||||
}
|
||||
|
||||
assert get_specific_rag_search_tool_config(user) == {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "dynamic",
|
||||
"tool_description": "Use this tool to search French legal documents and laws.",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_add_document_rag_search_tool_from_setting_adds_tools(settings):
|
||||
"""Test that add_document_rag_search_tool_from_setting adds tools to the agent."""
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
}
|
||||
|
||||
user = UserFactory()
|
||||
|
||||
agent = Agent("test")
|
||||
assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access
|
||||
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
# Check that tools were added
|
||||
assert len(agent._function_toolset.tools) == 1 # pylint: disable=protected-access
|
||||
assert agent._function_toolset.tools["legal_documents"].name == "legal_documents" # pylint: disable=protected-access
|
||||
assert (
|
||||
agent._function_toolset.tools["legal_documents"].description # pylint: disable=protected-access
|
||||
== "Use this tool to search legal documents and laws."
|
||||
)
|
||||
assert agent._function_toolset.tools["legal_documents"].function_schema.json_schema == { # pylint: disable=protected-access
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"query": {"description": "The query to search information about.", "type": "string"}
|
||||
},
|
||||
"required": ["query"],
|
||||
"type": "object",
|
||||
}
|
||||
|
||||
|
||||
def test_add_document_rag_search_tool_with_invalid_backend(settings, caplog):
|
||||
"""Test that invalid backend import is handled gracefully."""
|
||||
caplog.set_level(logging.WARNING, logger="chat.tools.document_generic_search_rag")
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"rag_backend_name": "non.existent.Backend",
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
}
|
||||
user = UserFactory()
|
||||
agent = Agent("test")
|
||||
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
# Tool should not be added due to import error
|
||||
assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access
|
||||
|
||||
# Check that warning was logged
|
||||
assert len(caplog.records) == 1
|
||||
assert "Could not import RAG backend non.existent.Backend" in caplog.records[0].message
|
||||
|
||||
|
||||
def test_add_document_rag_search_tool_with_missing_collection_ids(settings, caplog):
|
||||
"""Test that missing collection_ids is handled gracefully."""
|
||||
caplog.set_level(logging.WARNING, logger="chat.tools.document_generic_search_rag")
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
}
|
||||
user = UserFactory()
|
||||
agent = Agent("test")
|
||||
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
# Tool should not be added due to import error
|
||||
assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access
|
||||
|
||||
# Check that warning was logged
|
||||
assert len(caplog.records) == 1
|
||||
assert "No collection IDs provided for tool legal_documents" in caplog.records[0].message
|
||||
|
||||
|
||||
def test_add_document_rag_search_tool_with_missing_tool_description(settings, caplog):
|
||||
"""Test that missing tool_description is handled gracefully."""
|
||||
caplog.set_level(logging.WARNING, logger="chat.tools.document_generic_search_rag")
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
},
|
||||
}
|
||||
user = UserFactory()
|
||||
agent = Agent("test")
|
||||
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
# Tool should not be added due to import error
|
||||
assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access
|
||||
|
||||
# Check that warning was logged
|
||||
assert len(caplog.records) == 1
|
||||
assert "No tool description provided for tool legal_documents" in caplog.records[0].message
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_document_search_rag_tool_execution(settings):
|
||||
"""Test that the generated RAG tool executes correctly."""
|
||||
search_mock = respx.post("https://albert.api.etalab.gouv.fr/v1/search").mock(
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
json={
|
||||
"data": [
|
||||
{
|
||||
"method": "semantic",
|
||||
"chunk": {
|
||||
"id": 1,
|
||||
"content": "Relevant content snippet.",
|
||||
"metadata": {"document_name": "doc1.txt"},
|
||||
},
|
||||
"score": 0.9,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20},
|
||||
},
|
||||
)
|
||||
)
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
"legal_documents_2": {
|
||||
"collection_ids": [200],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
}
|
||||
user = UserFactory()
|
||||
agent = Agent(model="test")
|
||||
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
result = agent.run_sync("What information can you find about French services?")
|
||||
|
||||
# Verify the result
|
||||
assert json.loads(result.output) == {
|
||||
"legal_documents": {"0": {"snippets": "Relevant content snippet.", "url": "doc1.txt"}},
|
||||
"legal_documents_2": {"0": {"snippets": "Relevant content snippet.", "url": "doc1.txt"}},
|
||||
}
|
||||
|
||||
assert len(search_mock.calls) == 2
|
||||
assert json.loads(search_mock.calls[0].request.content) == {
|
||||
"collections": [100, 101, 102],
|
||||
"k": 4,
|
||||
"prompt": "a",
|
||||
"score_threshold": 0.6,
|
||||
}
|
||||
assert json.loads(search_mock.calls[1].request.content) == {
|
||||
"collections": [200],
|
||||
"k": 4,
|
||||
"prompt": "a",
|
||||
"score_threshold": 0.6,
|
||||
}
|
||||
|
||||
|
||||
def test_get_specific_rag_search_tool_config_with_empty_settings(settings):
|
||||
"""Test get_specific_rag_search_tool_config with empty SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS."""
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {}
|
||||
|
||||
user = UserFactory()
|
||||
config = get_specific_rag_search_tool_config(user)
|
||||
|
||||
assert config == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_add_document_rag_search_tool_function_call(settings):
|
||||
"""Test the function behavior."""
|
||||
search_mock = respx.post("https://albert.api.etalab.gouv.fr/v1/search").mock(
|
||||
return_value=httpx.Response(
|
||||
status_code=200,
|
||||
json={
|
||||
"data": [
|
||||
{
|
||||
"method": "semantic",
|
||||
"chunk": {
|
||||
"id": 1,
|
||||
"content": "Relevant content snippet.",
|
||||
"metadata": {"document_name": "doc1.txt"},
|
||||
},
|
||||
"score": 0.9,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20},
|
||||
},
|
||||
)
|
||||
)
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
}
|
||||
|
||||
user = await sync_to_async(UserFactory)()
|
||||
|
||||
agent = Agent("test")
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
result = await agent._function_toolset.tools["legal_documents"].function( # pylint: disable=protected-access
|
||||
RunContext(model="test", usage=RunUsage(), deps={}),
|
||||
query="Find information about French laws.",
|
||||
)
|
||||
|
||||
assert result.return_value == {
|
||||
"0": {"snippets": "Relevant content snippet.", "url": "doc1.txt"}
|
||||
}
|
||||
assert result.metadata == {"sources": {"doc1.txt"}}
|
||||
assert len(search_mock.calls) == 1
|
||||
assert json.loads(search_mock.calls[0].request.content) == {
|
||||
"collections": [100, 101, 102],
|
||||
"k": 4,
|
||||
"prompt": "Find information about French laws.",
|
||||
"score_threshold": 0.6,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_document_search_rag_http_status_error(settings, caplog):
|
||||
"""Test that HTTPStatusError is properly handled and logged."""
|
||||
caplog.set_level(logging.ERROR, logger="chat.tools.document_generic_search_rag")
|
||||
|
||||
# Mock the API to return a 500 error
|
||||
respx.post("https://albert.api.etalab.gouv.fr/v1/search").mock(
|
||||
return_value=httpx.Response(
|
||||
status_code=500,
|
||||
json={"error": "Internal server error"},
|
||||
)
|
||||
)
|
||||
|
||||
settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"legal_documents": {
|
||||
"collection_ids": [100, 101, 102],
|
||||
"feature_flag_value": "enabled",
|
||||
"tool_description": "Use this tool to search legal documents and laws.",
|
||||
},
|
||||
}
|
||||
|
||||
user = await sync_to_async(UserFactory)()
|
||||
agent = Agent("test")
|
||||
add_document_rag_search_tool_from_setting(agent, user)
|
||||
|
||||
# Call the tool function and expect a ModelRetry to be raised and caught
|
||||
tool_result = await agent._function_toolset.tools["legal_documents"].function( # pylint: disable=protected-access
|
||||
RunContext(model="test", usage=RunUsage(), deps={}),
|
||||
query="Find information about French laws.",
|
||||
)
|
||||
|
||||
# Verify the exception message
|
||||
assert tool_result == (
|
||||
"Document search service is currently unavailable: Server error '500 Internal "
|
||||
"Server Error' for url 'https://albert.api.etalab.gouv.fr/v1/search'\n"
|
||||
"For more information check: "
|
||||
"https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/500 You must "
|
||||
"explain this to the user and not try to answer based on your knowledge."
|
||||
)
|
||||
|
||||
# Verify that error was logged
|
||||
assert "RAG document search failed for tool legal_documents" in caplog.records[0].message
|
||||
assert "Document search service is currently unavailable" in caplog.records[1].message
|
||||
@@ -1,472 +0,0 @@
|
||||
"""Tests for document_summarize functionality."""
|
||||
|
||||
import io
|
||||
from unittest import mock
|
||||
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import ModelResponse, RunContext, TextPart
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.models.function import FunctionModel
|
||||
from pydantic_ai.usage import RunUsage
|
||||
|
||||
from chat.agents.summarize import SummarizationAgent
|
||||
from chat.llm_configuration import LLModel, LLMProvider
|
||||
from chat.tools.document_summarize import document_summarize, summarize_chunk
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fixture_summarization_agent_config(settings):
|
||||
"""Fixture to set used settings for agent configuration."""
|
||||
settings.LLM_CONFIGURATIONS = {
|
||||
settings.LLM_SUMMARIZATION_MODEL_HRID: LLModel(
|
||||
hrid="mistral-model",
|
||||
model_name="mistral-7b-instruct-v0.1",
|
||||
human_readable_name="Mistral 7B Instruct",
|
||||
profile=None,
|
||||
provider=LLMProvider(
|
||||
hrid="mistral",
|
||||
kind="mistral",
|
||||
base_url="https://api.mistral.ai/v1",
|
||||
api_key="testkey",
|
||||
),
|
||||
is_active=True,
|
||||
system_prompt="direct",
|
||||
tools=[],
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(name="mocked_context")
|
||||
def fixture_mocked_context():
|
||||
"""Fixture for a mocked RunContext."""
|
||||
mock_ctx = mock.Mock(spec=RunContext)
|
||||
mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0)
|
||||
mock_ctx.max_retries = 2
|
||||
mock_ctx.retries = {}
|
||||
return mock_ctx
|
||||
|
||||
|
||||
def mocked_summary(_messages, _info=None):
|
||||
"""Mocked summary response."""
|
||||
return ModelResponse(parts=[TextPart(content="This is a summary of the test chunk.")])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarize_chunk_returns_summary(mocked_context):
|
||||
"""Test that summarize_chunk returns a summary."""
|
||||
summarization_agent = SummarizationAgent()
|
||||
|
||||
with summarization_agent.override(model=FunctionModel(mocked_summary)):
|
||||
chunk = "This is a test chunk of text that needs to be summarized."
|
||||
|
||||
result = await summarize_chunk(1, chunk, 1, summarization_agent, mocked_context)
|
||||
|
||||
assert result == "This is a summary of the test chunk."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarize_chunk_raises_model_retry_on_error(mocked_context):
|
||||
"""Test that summarize_chunk raises ModelRetry when agent fails."""
|
||||
summarization_agent = SummarizationAgent()
|
||||
|
||||
def mocked_summary_error(_messages, _info=None):
|
||||
"""Mocked summary that raises an error."""
|
||||
raise ValueError("Simulated LLM error")
|
||||
|
||||
with summarization_agent.override(model=FunctionModel(mocked_summary_error)):
|
||||
chunk = "This is a test chunk."
|
||||
|
||||
with pytest.raises(ModelRetry) as exc_info:
|
||||
await summarize_chunk(1, chunk, 1, summarization_agent, mocked_context)
|
||||
|
||||
assert "An error occurred while summarizing a part of the document chunk" in str(
|
||||
exc_info.value
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_summarize_chunk_handles_empty_response(mocked_context):
|
||||
"""Test that summarize_chunk handles empty responses from the agent."""
|
||||
summarization_agent = SummarizationAgent()
|
||||
|
||||
def mocked_empty_summary(_messages, _info=None):
|
||||
"""Mocked summary that returns empty content."""
|
||||
return ModelResponse(parts=[TextPart(content="")])
|
||||
|
||||
with summarization_agent.override(model=FunctionModel(mocked_empty_summary)):
|
||||
chunk = "This is a test chunk."
|
||||
|
||||
# Empty responses cause ModelRetry since pydantic-ai considers them invalid
|
||||
with pytest.raises(ModelRetry):
|
||||
await summarize_chunk(1, chunk, 1, summarization_agent, mocked_context)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_single_document(
|
||||
settings, mocked_context, mock_summarization_agent
|
||||
):
|
||||
"""Test document_summarize with a single document."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 100
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 10
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
# Create mock conversation with a text attachment
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "test_doc.txt"
|
||||
mock_attachment.file_name = "test_doc.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
# Mock file storage
|
||||
file_content = "This is a test document. " * 20 # Create a document with some content
|
||||
with mock.patch.object(
|
||||
default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8"))
|
||||
):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
call_count = {"chunk": 0, "merge": 0}
|
||||
|
||||
def mocked_summary_full(messages, _info=None):
|
||||
"""Mocked summary for full flow."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" in messages_text:
|
||||
call_count["merge"] += 1
|
||||
return ModelResponse(
|
||||
parts=[TextPart(content="# Final Summary\n\nThis is the final merged summary.")]
|
||||
)
|
||||
|
||||
call_count["chunk"] += 1
|
||||
return ModelResponse(
|
||||
parts=[TextPart(content=f"Summary of chunk {call_count['chunk']}")]
|
||||
)
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary_full)):
|
||||
result = await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
assert result.return_value == "# Final Summary\n\nThis is the final merged summary."
|
||||
assert result.metadata["sources"] == {"test_doc.txt"}
|
||||
assert call_count["merge"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_multiple_documents(
|
||||
settings, mocked_context, mock_summarization_agent
|
||||
):
|
||||
"""Test document_summarize with multiple documents."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 50
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 5
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
# Create mock conversation with multiple text attachments
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment1 = mock.Mock()
|
||||
mock_attachment1.key = "doc1.txt"
|
||||
mock_attachment1.file_name = "doc1.txt"
|
||||
mock_attachment1.content_type = "text/plain"
|
||||
|
||||
mock_attachment2 = mock.Mock()
|
||||
mock_attachment2.key = "doc2.txt"
|
||||
mock_attachment2.file_name = "doc2.txt"
|
||||
mock_attachment2.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment1, mock_attachment2]
|
||||
|
||||
file_content1 = "Content of document one. " * 10
|
||||
file_content2 = "Content of document two. " * 10
|
||||
|
||||
def mock_open_side_effect(key):
|
||||
"""Mock file opening based on key."""
|
||||
if key == "doc1.txt":
|
||||
return io.BytesIO(file_content1.encode("utf-8"))
|
||||
return io.BytesIO(file_content2.encode("utf-8"))
|
||||
|
||||
with mock.patch.object(default_storage, "open", side_effect=mock_open_side_effect):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
def mocked_summary_multi(messages, _info=None):
|
||||
"""Mocked summary for multiple documents."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" in messages_text:
|
||||
return ModelResponse(parts=[TextPart(content="Combined summary of all documents")])
|
||||
|
||||
return ModelResponse(parts=[TextPart(content="Chunk summary")])
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary_multi)):
|
||||
result = await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
assert result.return_value == "Combined summary of all documents"
|
||||
assert result.metadata["sources"] == {"doc1.txt", "doc2.txt"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_with_custom_instructions(
|
||||
settings, mocked_context, mock_summarization_agent
|
||||
):
|
||||
"""Test document_summarize with custom instructions."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 100
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 10
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "test.txt"
|
||||
mock_attachment.file_name = "test.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
file_content = "Test content " * 20
|
||||
|
||||
with mock.patch.object(
|
||||
default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8"))
|
||||
):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
captured_merge_prompt = []
|
||||
|
||||
def mocked_summary_with_instructions(messages, _info=None):
|
||||
"""Mocked summary that captures merge prompt."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" in messages_text:
|
||||
captured_merge_prompt.append(messages_text)
|
||||
return ModelResponse(parts=[TextPart(content="Summary in 2 paragraphs")])
|
||||
|
||||
return ModelResponse(parts=[TextPart(content="Chunk summary")])
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary_with_instructions)):
|
||||
result = await document_summarize(
|
||||
mocked_context, instructions="summary in 2 paragraphs"
|
||||
)
|
||||
|
||||
assert result.return_value == "Summary in 2 paragraphs"
|
||||
assert len(captured_merge_prompt) == 1
|
||||
assert "summary in 2 paragraphs" in captured_merge_prompt[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_no_text_attachments(mocked_context, mock_summarization_agent):
|
||||
"""Test document_summarize returns error message when no text documents found."""
|
||||
mock_conversation = mock.Mock()
|
||||
mock_conversation.attachments.filter.return_value = []
|
||||
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
# The decorator @last_model_retry_soft_fail catches ModelCannotRetry and returns a message
|
||||
# We still need to provide a mock agent even if it won't be called
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary)):
|
||||
result = await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
assert "No text documents found in the conversation" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_error_reading_document(mocked_context, mock_summarization_agent):
|
||||
"""Test document_summarize handles errors when reading documents."""
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "test.txt"
|
||||
mock_attachment.file_name = "test.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
with mock.patch.object(default_storage, "open", side_effect=IOError("File read error")):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
# The decorator @last_model_retry_soft_fail catches ModelCannotRetry and returns a message
|
||||
# We still need to provide a mock agent even if it won't be called
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary)):
|
||||
result = await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
assert "An unexpected error occurred during document summarization" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_error_during_chunk_summarization(
|
||||
settings, mocked_context, mock_summarization_agent
|
||||
):
|
||||
"""Test document_summarize handles ModelRetry during chunk summarization."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 100
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 10
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "test.txt"
|
||||
mock_attachment.file_name = "test.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
file_content = "Test content " * 20
|
||||
|
||||
with mock.patch.object(
|
||||
default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8"))
|
||||
):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
def mocked_summary_error(messages, _info=None):
|
||||
"""Mocked summary that raises an error during chunks."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" not in messages_text:
|
||||
raise ValueError("Chunk processing error")
|
||||
|
||||
return ModelResponse(parts=[TextPart(content="Final summary")])
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary_error)):
|
||||
with pytest.raises(ModelRetry):
|
||||
await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_error_during_merge(
|
||||
settings, mocked_context, mock_summarization_agent
|
||||
):
|
||||
"""Test document_summarize handles errors during final merge."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 100
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 10
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "test.txt"
|
||||
mock_attachment.file_name = "test.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
file_content = "Test content " * 20
|
||||
|
||||
with mock.patch.object(
|
||||
default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8"))
|
||||
):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
def mocked_summary_merge_error(messages, _info=None):
|
||||
"""Mocked summary that raises an error during merge."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" in messages_text:
|
||||
raise ValueError("Merge error")
|
||||
|
||||
return ModelResponse(parts=[TextPart(content="Chunk summary")])
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary_merge_error)):
|
||||
with pytest.raises(ModelRetry) as exc_info:
|
||||
await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
# Should raise ModelRetry regardless of which phase failed
|
||||
assert "An error occurred" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_empty_result(settings, mocked_context, mock_summarization_agent):
|
||||
"""Test document_summarize raises ModelRetry when summarization produces empty result."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 100
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 10
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "test.txt"
|
||||
mock_attachment.file_name = "test.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
file_content = "Test content " * 20
|
||||
|
||||
with mock.patch.object(
|
||||
default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8"))
|
||||
):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
def mocked_empty_summary(messages, _info=None):
|
||||
"""Mocked summary that returns empty for merge."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" in messages_text:
|
||||
return ModelResponse(parts=[TextPart(content=" ")])
|
||||
|
||||
return ModelResponse(parts=[TextPart(content="Chunk summary")])
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_empty_summary)):
|
||||
with pytest.raises(ModelRetry) as exc_info:
|
||||
await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
# Should raise ModelRetry with the specific message
|
||||
assert "The summarization produced an empty result" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_summarize_large_document_multiple_chunks(
|
||||
settings, mocked_context, mock_summarization_agent
|
||||
):
|
||||
"""Test document_summarize with a large document requiring multiple chunks."""
|
||||
settings.SUMMARIZATION_CHUNK_SIZE = 20 # Small chunk size to force multiple chunks
|
||||
settings.SUMMARIZATION_OVERLAP_SIZE = 5
|
||||
settings.SUMMARIZATION_CONCURRENT_REQUESTS = 2
|
||||
|
||||
mock_conversation = mock.Mock()
|
||||
mock_attachment = mock.Mock()
|
||||
mock_attachment.key = "large_doc.txt"
|
||||
mock_attachment.file_name = "large_doc.txt"
|
||||
mock_attachment.content_type = "text/plain"
|
||||
|
||||
mock_conversation.attachments.filter.return_value = [mock_attachment]
|
||||
|
||||
# Create a large document
|
||||
file_content = "This is a word. " * 100 # Should create multiple chunks
|
||||
|
||||
with mock.patch.object(
|
||||
default_storage, "open", return_value=io.BytesIO(file_content.encode("utf-8"))
|
||||
):
|
||||
# Set up mocked_context with conversation
|
||||
mocked_context.deps = mock.Mock()
|
||||
mocked_context.deps.conversation = mock_conversation
|
||||
|
||||
chunk_count = {"count": 0}
|
||||
|
||||
def mocked_summary_multi_chunks(messages, _info=None):
|
||||
"""Mocked summary that counts chunks."""
|
||||
messages_text = messages[0].parts[-1].content
|
||||
|
||||
if "Produce a coherent synthesis" in messages_text:
|
||||
return ModelResponse(
|
||||
parts=[TextPart(content=f"Final summary of {chunk_count['count']} chunks")]
|
||||
)
|
||||
|
||||
chunk_count["count"] += 1
|
||||
return ModelResponse(
|
||||
parts=[TextPart(content=f"Summary of chunk {chunk_count['count']}")]
|
||||
)
|
||||
|
||||
with mock_summarization_agent(FunctionModel(mocked_summary_multi_chunks)):
|
||||
result = await document_summarize(mocked_context, instructions=None)
|
||||
|
||||
assert "Final summary of" in result.return_value
|
||||
assert chunk_count["count"] > 1 # Ensure multiple chunks were processed
|
||||
@@ -1,154 +0,0 @@
|
||||
"""Tests for chat tool utilities."""
|
||||
|
||||
import inspect
|
||||
from typing import get_type_hints
|
||||
|
||||
import pytest
|
||||
from pydantic_ai import ModelRetry, RunContext
|
||||
|
||||
from chat.tools.exceptions import ModelCannotRetry
|
||||
from chat.tools.utils import last_model_retry_soft_fail
|
||||
|
||||
|
||||
def test_last_model_retry_soft_fail_preserves_function_metadata():
|
||||
"""Test that the decorator preserves function metadata for schema generation."""
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def example_tool(ctx: RunContext, query: str, limit: int = 10) -> str: # pylint: disable=unused-argument
|
||||
"""
|
||||
Example tool function.
|
||||
|
||||
Args:
|
||||
ctx: The run context.
|
||||
query: The search query.
|
||||
limit: Maximum number of results.
|
||||
|
||||
Returns:
|
||||
The search results.
|
||||
"""
|
||||
return f"Results for {query} (limit: {limit})"
|
||||
|
||||
# Check that function name is preserved
|
||||
assert example_tool.__name__ == "example_tool"
|
||||
|
||||
# Check that docstring is preserved
|
||||
assert example_tool.__doc__ is not None
|
||||
assert "Example tool function" in example_tool.__doc__
|
||||
|
||||
# Check that signature is preserved
|
||||
sig = inspect.signature(example_tool)
|
||||
assert "ctx" in sig.parameters
|
||||
assert "query" in sig.parameters
|
||||
assert "limit" in sig.parameters
|
||||
assert sig.parameters["limit"].default == 10
|
||||
|
||||
# Check that type hints are preserved
|
||||
type_hints = get_type_hints(example_tool)
|
||||
assert "query" in type_hints
|
||||
assert type_hints["query"] == str
|
||||
assert "limit" in type_hints
|
||||
assert type_hints["limit"] == int
|
||||
assert type_hints["return"] == str
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_last_model_retry_soft_fail_normal_execution():
|
||||
"""Test that the decorator doesn't interfere with normal execution."""
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def example_tool(_ctx: RunContext, value: str) -> str:
|
||||
"""Example tool."""
|
||||
return f"Result: {value}"
|
||||
|
||||
# Create a mock context
|
||||
class MockContext:
|
||||
"""Fake context for testing."""
|
||||
|
||||
max_retries = 3
|
||||
retries = {}
|
||||
tool_name = "example_tool"
|
||||
|
||||
ctx = MockContext()
|
||||
result = await example_tool(ctx, "test")
|
||||
assert result == "Result: test"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_last_model_retry_soft_fail_handles_retry_exception():
|
||||
"""Test that the decorator handles ModelRetry exceptions correctly."""
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def failing_tool(_ctx: RunContext, should_fail: bool) -> str:
|
||||
"""Tool that can raise ModelRetry."""
|
||||
if should_fail:
|
||||
raise ModelRetry("Please retry with different parameters")
|
||||
return "Success"
|
||||
|
||||
# Create a mock context
|
||||
class MockContext:
|
||||
"""Fake context for testing."""
|
||||
|
||||
max_retries = 3
|
||||
retries = {}
|
||||
tool_name = "failing_tool"
|
||||
|
||||
ctx = MockContext()
|
||||
|
||||
# Test when retries haven't been exhausted - should re-raise
|
||||
with pytest.raises(ModelRetry):
|
||||
await failing_tool(ctx, should_fail=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_last_model_retry_soft_fail_returns_message_when_max_retries_reached():
|
||||
"""Test that the decorator returns the error message when max retries is reached."""
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def failing_tool(_ctx: RunContext, should_fail: bool) -> str:
|
||||
"""Tool that can raise ModelRetry."""
|
||||
if should_fail:
|
||||
raise ModelRetry("Please retry with different parameters.")
|
||||
return "Success"
|
||||
|
||||
# Create a mock context with max retries already reached
|
||||
class MockContext:
|
||||
"""Fake context for testing."""
|
||||
|
||||
max_retries = 3
|
||||
retries = {"failing_tool": 3}
|
||||
tool_name = "failing_tool"
|
||||
|
||||
ctx = MockContext()
|
||||
|
||||
# Test when retries have been exhausted - should return message
|
||||
result = await failing_tool(ctx, should_fail=True)
|
||||
assert result == (
|
||||
"Please retry with different parameters. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_last_model_retry_soft_fail_returns_message_when_model_cannot_retry():
|
||||
"""Test that the decorator returns the error message when ModelCannotRetry is raised."""
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def failing_tool(_ctx: RunContext, should_fail: bool) -> str:
|
||||
"""Tool that can raise ModelRetry."""
|
||||
if should_fail:
|
||||
raise ModelCannotRetry("This is broken duh.")
|
||||
return "Success"
|
||||
|
||||
# Create a mock context with max retries already reached
|
||||
class MockContext:
|
||||
"""Fake context for testing."""
|
||||
|
||||
max_retries = 3
|
||||
retries = {"failing_tool": 3}
|
||||
tool_name = "failing_tool"
|
||||
|
||||
ctx = MockContext()
|
||||
|
||||
# Test when retries have been exhausted - should return message
|
||||
result = await failing_tool(ctx, should_fail=True)
|
||||
assert result == "This is broken duh."
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,11 +0,0 @@
|
||||
"""tools for testing chat functionality"""
|
||||
|
||||
import re
|
||||
|
||||
|
||||
def replace_uuids_with_placeholder(text):
|
||||
"""Replace all UUIDs in the given text with a placeholder."""
|
||||
text = re.sub('"toolCallId":"([a-z0-9-]){36}"', '"toolCallId":"XXX"', text)
|
||||
text = re.sub('"toolCallId":"pyd_ai_([a-z0-9]){32}"', '"toolCallId":"pyd_ai_YYY"', text)
|
||||
text = re.sub('"([a-z0-9-]){36}"', '"<mocked_uuid>"', text)
|
||||
return text
|
||||
@@ -1,268 +0,0 @@
|
||||
"""Tests for chat attachment views."""
|
||||
|
||||
import uuid
|
||||
from io import BytesIO
|
||||
from unittest import mock
|
||||
|
||||
from django.core.files.storage import default_storage
|
||||
from django.test import override_settings
|
||||
|
||||
import pytest
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
|
||||
from chat import factories, models
|
||||
from chat.tests.conftest import PIXEL_PNG
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
def test_attachment_create_anonymous_forbidden(api_client):
|
||||
"""
|
||||
Anonymous users should not be able to create attachments.
|
||||
"""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/"
|
||||
response = api_client.post(url, {"file_name": "test.png", "size": 123}, format="json")
|
||||
|
||||
assert response.status_code == 401
|
||||
assert response.json() == {"detail": "Authentication credentials were not provided."}
|
||||
|
||||
|
||||
def test_attachment_create_authenticated_not_owner_forbidden(api_client):
|
||||
"""
|
||||
A user who does not own the conversation should not be able to create an attachment.
|
||||
"""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
user = factories.UserFactory()
|
||||
api_client.force_login(user)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/"
|
||||
response = api_client.post(
|
||||
url, {"file_name": "test.png", "size": 123, "content_type": "image/png"}, format="json"
|
||||
)
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
def test_attachment_create_success(api_client):
|
||||
"""
|
||||
An authenticated user who owns the conversation should be able to create an attachment.
|
||||
"""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/"
|
||||
|
||||
response = api_client.post(
|
||||
url, {"file_name": "test.png", "size": 123, "content_type": "image/png"}, format="json"
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
data = response.json()
|
||||
assert data["policy"] is not None
|
||||
assert data["key"].startswith(f"{conversation.pk!s}/attachments/")
|
||||
assert data["key"].endswith(".png")
|
||||
|
||||
attachment = models.ChatConversationAttachment.objects.get(pk=data["id"])
|
||||
assert attachment.conversation == conversation
|
||||
assert attachment.uploaded_by == conversation.owner
|
||||
assert attachment.upload_state == AttachmentStatus.PENDING
|
||||
assert attachment.file_name == "test.png"
|
||||
assert attachment.size == 123
|
||||
assert attachment.content_type == "image/png"
|
||||
|
||||
|
||||
def test_attachment_create_size_limit_exceeded(api_client, settings):
|
||||
"""
|
||||
The attachment should not be created if the file size exceeds the maximum limit.
|
||||
"""
|
||||
settings.ATTACHMENT_MAX_SIZE = 1024 # 1 KB for test
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/"
|
||||
response = api_client.post(
|
||||
url, {"file_name": "test.png", "size": 2048, "content_type": "image/png"}, format="json"
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.json() == {"size": ["File size exceeds the maximum limit of 0 MB."]}
|
||||
|
||||
|
||||
def test_attachment_retrieve_anonymous_forbidden(api_client):
|
||||
"""
|
||||
Anonymous users should not be able to retrieve attachments.
|
||||
"""
|
||||
attachment = factories.ChatConversationAttachmentFactory()
|
||||
url = f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/"
|
||||
response = api_client.get(url)
|
||||
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
def test_attachment_retrieve_not_owner_forbidden(api_client):
|
||||
"""
|
||||
A user who does not own the conversation should not be able to retrieve an attachment.
|
||||
"""
|
||||
attachment = factories.ChatConversationAttachmentFactory()
|
||||
user = factories.UserFactory()
|
||||
api_client.force_login(user)
|
||||
|
||||
url = f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/"
|
||||
response = api_client.get(url)
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"upload_state, expected_url_present",
|
||||
[
|
||||
(AttachmentStatus.PENDING, False),
|
||||
(AttachmentStatus.ANALYZING, False),
|
||||
(AttachmentStatus.READY, True),
|
||||
(AttachmentStatus.SUSPICIOUS, True),
|
||||
(AttachmentStatus.FILE_TOO_LARGE_TO_ANALYZE, True),
|
||||
],
|
||||
)
|
||||
def test_attachment_retrieve_success(api_client, upload_state, expected_url_present):
|
||||
"""
|
||||
An authenticated user who owns the conversation should be able to retrieve an attachment.
|
||||
The URL should be present only for appropriate statuses.
|
||||
"""
|
||||
attachment = factories.ChatConversationAttachmentFactory(upload_state=upload_state)
|
||||
api_client.force_login(attachment.conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/"
|
||||
response = api_client.get(url)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["id"] == str(attachment.pk)
|
||||
assert ("url" in data and data["url"] is not None) == expected_url_present
|
||||
|
||||
|
||||
@override_settings(POSTHOG_KEY="test_key")
|
||||
@mock.patch("chat.views.posthog")
|
||||
@mock.patch("chat.views.malware_detection.analyse_file")
|
||||
def test_upload_ended_success(mock_analyse_file, mock_posthog, api_client):
|
||||
"""
|
||||
The 'upload_ended' action should change the attachment state and trigger analysis.
|
||||
"""
|
||||
attachment = factories.ChatConversationAttachmentFactory(
|
||||
upload_state=AttachmentStatus.PENDING,
|
||||
file_name="test.txt",
|
||||
size=4,
|
||||
)
|
||||
api_client.force_login(attachment.conversation.owner)
|
||||
|
||||
# Create a dummy file in storage
|
||||
default_storage.connection.meta.client.put_object(
|
||||
Bucket=default_storage.bucket_name,
|
||||
Key=attachment.key,
|
||||
Body=BytesIO(b"my prose"),
|
||||
ContentType="text/plain",
|
||||
)
|
||||
|
||||
url = (
|
||||
f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/upload-ended/"
|
||||
)
|
||||
response = api_client.post(url)
|
||||
|
||||
assert response.status_code == 200
|
||||
attachment.refresh_from_db()
|
||||
assert attachment.upload_state == AttachmentStatus.ANALYZING
|
||||
assert attachment.content_type == "text/plain"
|
||||
|
||||
mock_analyse_file.assert_called_once_with(
|
||||
attachment.key,
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=attachment.conversation.pk,
|
||||
)
|
||||
mock_posthog.capture.assert_called_once()
|
||||
|
||||
|
||||
def test_upload_ended_not_pending(api_client):
|
||||
"""
|
||||
The 'upload_ended' action should fail if the attachment is not in the PENDING state.
|
||||
"""
|
||||
attachment = factories.ChatConversationAttachmentFactory(
|
||||
upload_state=AttachmentStatus.ANALYZING
|
||||
)
|
||||
api_client.force_login(attachment.conversation.owner)
|
||||
|
||||
url = (
|
||||
f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/upload-ended/"
|
||||
)
|
||||
response = api_client.post(url)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert response.json() == {
|
||||
"attachment": "This action is only available for items in PENDING state."
|
||||
}
|
||||
|
||||
|
||||
def test_upload_ended_not_owner(api_client):
|
||||
"""
|
||||
A user who does not own the conversation cannot end the upload.
|
||||
"""
|
||||
attachment = factories.ChatConversationAttachmentFactory(upload_state=AttachmentStatus.PENDING)
|
||||
user = factories.UserFactory()
|
||||
api_client.force_login(user)
|
||||
|
||||
url = (
|
||||
f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/upload-ended/"
|
||||
)
|
||||
response = api_client.post(url)
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name,content,_extension,content_type",
|
||||
[
|
||||
("test.exe", b"text", "exe", "text/plain"),
|
||||
("test", b"text", "txt", "text/plain"),
|
||||
("test.aaaaaa", b"test", "txt", "text/plain"),
|
||||
("test.txt", PIXEL_PNG, "txt", "image/png"),
|
||||
("test.py", b"#!/usr/bin/python", "py", "text/plain"),
|
||||
],
|
||||
)
|
||||
def test_upload_ended_fix_extension(api_client, name, content, _extension, content_type):
|
||||
"""
|
||||
The 'upload_ended' action should update the attachment's file_name, content_type, and size
|
||||
based on the actual uploaded file.
|
||||
"""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
attachment = factories.ChatConversationAttachmentFactory(
|
||||
conversation=conversation,
|
||||
upload_state=AttachmentStatus.PENDING,
|
||||
file_name=name,
|
||||
size=0,
|
||||
key=f"{conversation.pk!s}/attachments/temp-{uuid.uuid4()!s}-{name:s}",
|
||||
content_type="application/wrong",
|
||||
)
|
||||
|
||||
# Create a dummy file in storage
|
||||
default_storage.connection.meta.client.put_object(
|
||||
Bucket=default_storage.bucket_name,
|
||||
Key=attachment.key,
|
||||
Body=content,
|
||||
ContentType=content_type,
|
||||
)
|
||||
|
||||
api_client.force_login(conversation.owner)
|
||||
url = (
|
||||
f"/api/v1.0/chats/{attachment.conversation.pk}/attachments/{attachment.pk!s}/upload-ended/"
|
||||
)
|
||||
response = api_client.post(url)
|
||||
|
||||
assert response.status_code == 200
|
||||
attachment.refresh_from_db()
|
||||
|
||||
assert attachment.upload_state == AttachmentStatus.READY # malware_detection mocked to safe
|
||||
assert attachment.content_type == content_type # updated
|
||||
assert attachment.file_name == name # updated
|
||||
assert attachment.size == len(content) # updated
|
||||
@@ -1,6 +1,8 @@
|
||||
"""Common test fixtures for chat conversation endpoint tests."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
@@ -10,19 +12,23 @@ import respx
|
||||
from freezegun import freeze_time
|
||||
|
||||
|
||||
def build_openai_stream():
|
||||
@pytest.fixture(name="mock_uuid4")
|
||||
def mock_uuid4_fixture():
|
||||
"""Fixture to mock UUID generation for testing."""
|
||||
value = uuid.uuid4()
|
||||
with patch("uuid.uuid4", return_value=value):
|
||||
yield value
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_openai_stream")
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def fixture_mock_openai_stream():
|
||||
"""
|
||||
Constructs a string that simulates an OpenAI streaming response payload.
|
||||
|
||||
The returned string contains three OpenAI-style `data:` blocks: a first chunk with content "Hello",
|
||||
a second chunk with content " there" and a `finish_reason` of "stop" (including a `usage` object),
|
||||
and a final `data: [DONE]` marker. Timestamp fields are generated from timezone.now() converted to
|
||||
naive timestamps.
|
||||
|
||||
Returns:
|
||||
A string containing concatenated `data:` lines representing streaming chunks and a final `[DONE]` marker.
|
||||
Fixture to mock the OpenAI stream response.
|
||||
|
||||
See https://platform.openai.com/docs/api-reference/chat-streaming/streaming
|
||||
"""
|
||||
return (
|
||||
openai_stream = (
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
@@ -63,24 +69,7 @@ def build_openai_stream():
|
||||
"data: [DONE]\n\n"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_openai_stream")
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def fixture_mock_openai_stream():
|
||||
"""
|
||||
Fixture to mock the OpenAI stream response.
|
||||
|
||||
See https://platform.openai.com/docs/api-reference/chat-streaming/streaming
|
||||
"""
|
||||
openai_stream = build_openai_stream()
|
||||
|
||||
async def mock_stream():
|
||||
"""
|
||||
Yield each line of the prepared OpenAI-style streaming payload as encoded bytes.
|
||||
|
||||
Yields:
|
||||
AsyncGenerator[bytes, None]: Sequential byte chunks for each line in the constructed stream, preserving original line endings.
|
||||
"""
|
||||
for line in openai_stream.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
@@ -91,100 +80,10 @@ def fixture_mock_openai_stream():
|
||||
return route
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_openai_stream_with_title_generation")
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def fixture_mock_openai_stream_with_title_generation():
|
||||
"""
|
||||
Mock pytest fixture that intercepts POST requests to the external chat completions endpoint and returns either a streaming chat response or a non-streaming title-generation response depending on the incoming request.
|
||||
|
||||
When the request JSON has "stream" set to True, the fixture returns an HTTP streaming response that imitates OpenAI's chat streaming payload; otherwise it returns a non-streaming JSON response containing a generated title and usage metadata.
|
||||
|
||||
Returns:
|
||||
respx.Route: A configured respx route that intercepts POST requests to
|
||||
"https://www.external-ai-service.com/chat/completions" and replies based on the request body.
|
||||
"""
|
||||
|
||||
def create_stream_response():
|
||||
"""
|
||||
Create an HTTP response whose body streams encoded lines of an OpenAI-style streaming payload.
|
||||
|
||||
Returns:
|
||||
httpx.Response: HTTP 200 response with a streaming body that yields encoded bytes for each line of the streaming payload.
|
||||
"""
|
||||
openai_stream = build_openai_stream()
|
||||
|
||||
async def mock_stream():
|
||||
"""
|
||||
Yield encoded byte chunks for each line of the OpenAI stream.
|
||||
|
||||
Each yielded value is a bytes object containing one line (including its line ending) from the prebuilt OpenAI streaming payload, suitable for use as an HTTP streaming response body.
|
||||
"""
|
||||
for line in openai_stream.splitlines(keepends=True):
|
||||
yield line.encode()
|
||||
|
||||
return httpx.Response(200, stream=mock_stream())
|
||||
|
||||
def create_non_stream_response():
|
||||
"""
|
||||
Create a non-streaming OpenAI-like chat completion response containing a generated title.
|
||||
|
||||
Returns:
|
||||
httpx.Response: HTTP 200 response whose JSON payload represents a chat completion with a single assistant message containing the generated title and accompanying metadata (id, model, timestamps, choices, and usage).
|
||||
"""
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-title",
|
||||
"object": "chat.completion",
|
||||
"created": int(timezone.make_naive(timezone.now()).timestamp()),
|
||||
"model": "test-model",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "GENERATED TITLE",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 50, "completion_tokens": 5, "total_tokens": 55},
|
||||
},
|
||||
)
|
||||
|
||||
def handle_request(request):
|
||||
"""
|
||||
Selects a streaming or non-streaming HTTP response based on the request JSON `stream` flag.
|
||||
|
||||
Parameters:
|
||||
request (httpx.Request): Incoming request whose JSON body is inspected for the `stream` boolean flag.
|
||||
|
||||
Returns:
|
||||
httpx.Response: A response that streams the OpenAI-style event lines if `stream` is True, otherwise a non-streaming JSON response.
|
||||
"""
|
||||
body = json.loads(request.content)
|
||||
if body.get("stream", False):
|
||||
return create_stream_response()
|
||||
return create_non_stream_response()
|
||||
|
||||
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
|
||||
side_effect=handle_request
|
||||
)
|
||||
|
||||
return route
|
||||
|
||||
|
||||
@pytest.fixture(name="mock_openai_no_stream")
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def fixture_mock_openai_no_stream():
|
||||
"""
|
||||
Create a respx route that returns a fixed, non-streaming OpenAI chat completion response.
|
||||
|
||||
The mocked response is an HTTP 200 JSON payload representing a completed assistant message (explaining Rayleigh scattering) with associated metadata and usage details.
|
||||
|
||||
Returns:
|
||||
respx.Route: The configured respx route intercepting POST requests to https://www.external-ai-service.com/chat/completions.
|
||||
"""
|
||||
"""Fixture to mock the OpenAI response."""
|
||||
|
||||
route = respx.post("https://www.external-ai-service.com/chat/completions").mock(
|
||||
return_value=httpx.Response(
|
||||
@@ -498,4 +397,4 @@ def fixture_mock_openai_stream_tool():
|
||||
]
|
||||
)
|
||||
|
||||
return route
|
||||
return route
|
||||
|
||||
@@ -2,14 +2,11 @@
|
||||
# pylint: disable=too-many-lines
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
from asgiref.sync import sync_to_async
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from rest_framework import status
|
||||
|
||||
@@ -24,7 +21,6 @@ from chat.ai_sdk_types import (
|
||||
)
|
||||
from chat.factories import ChatConversationFactory
|
||||
from chat.llm_configuration import LLModel, LLMProvider
|
||||
from chat.tests.utils import replace_uuids_with_placeholder
|
||||
|
||||
# enable database transactions for tests:
|
||||
# transaction=True ensures that the data are available in the database
|
||||
@@ -90,7 +86,7 @@ def test_post_conversation_invalid_protocol(api_client):
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
def test_post_conversation_data_protocol(api_client, mock_openai_stream, mock_uuid4):
|
||||
"""Test posting messages to a conversation using the 'data' protocol."""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
|
||||
@@ -117,14 +113,10 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"Hello"\n'
|
||||
'0:" there"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -143,9 +135,8 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -156,9 +147,8 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -169,7 +159,6 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
parts=[TextUIPart(type="text", text="Hello there")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -199,7 +188,6 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -220,14 +208,13 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
def test_post_conversation_text_protocol(api_client, mock_openai_stream, mock_uuid4):
|
||||
"""Test posting messages to a conversation using the 'text' protocol."""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
|
||||
@@ -269,9 +256,8 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -282,9 +268,8 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -295,7 +280,6 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
parts=[TextUIPart(type="text", text="Hello there")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -325,7 +309,6 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -346,14 +329,13 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock_uuid4):
|
||||
"""Ensure an image URL is correctly forwarded to the AI service."""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
url = f"/api/v1.0/chats/{chat_conversation.pk}/conversation/?protocol=data"
|
||||
@@ -391,14 +373,10 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"I see a cat"\n'
|
||||
'0:" in the picture."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -459,9 +437,8 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello, what do you see on this picture?",
|
||||
reasoning=None,
|
||||
@@ -482,9 +459,8 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
parts=[TextUIPart(type="text", text="Hello, what do you see on this picture?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I see a cat in the picture.",
|
||||
reasoning=None,
|
||||
@@ -495,7 +471,6 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
parts=[TextUIPart(type="text", text="I see a cat in the picture.")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -537,7 +512,6 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -558,14 +532,13 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settings):
|
||||
def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_uuid4, settings):
|
||||
"""Ensure tool calls are correctly forwarded and streamed back."""
|
||||
settings.AI_AGENT_TOOLS = ["get_current_weather"]
|
||||
|
||||
@@ -594,10 +567,6 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'b:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","toolName":'
|
||||
'"get_current_weather"}\n'
|
||||
@@ -606,7 +575,7 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
'a:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","result":{"location":'
|
||||
'"Paris","temperature":22,"unit":"celsius"}}\n'
|
||||
'0:"The current weather in Paris is nice"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -637,9 +606,8 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -650,9 +618,8 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
parts=[TextUIPart(type="text", text="Weather in Paris?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The current weather in Paris is nice",
|
||||
reasoning=None,
|
||||
@@ -675,7 +642,6 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -705,7 +671,6 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "tool_call",
|
||||
@@ -734,7 +699,6 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -749,7 +713,6 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
"tool_name": "get_current_weather",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -772,14 +735,15 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool, settings):
|
||||
def test_post_conversation_tool_call_fails(
|
||||
api_client, mock_openai_stream_tool, mock_uuid4, settings
|
||||
):
|
||||
"""Ensure tool calls are correctly forwarded and streamed back when failing."""
|
||||
settings.AI_AGENT_TOOLS = []
|
||||
|
||||
@@ -808,10 +772,6 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'b:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","toolName":"get_current_weather"}\n'
|
||||
'c:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","argsTextDelta":'
|
||||
@@ -819,7 +779,7 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
'a:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","result":"Unknown tool '
|
||||
"name: 'get_current_weather'. No tools available.\"}\n"
|
||||
'0:"I cannot give you an answer to that."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -850,9 +810,8 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -863,9 +822,8 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
parts=[TextUIPart(type="text", text="Weather in Paris?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I cannot give you an answer to that.",
|
||||
reasoning=None,
|
||||
@@ -888,7 +846,6 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -918,7 +875,6 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "tool_call",
|
||||
@@ -947,7 +903,6 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -961,7 +916,6 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
"tool_name": "get_current_weather",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -984,7 +938,6 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -1017,6 +970,7 @@ def test_post_conversation_model_selection_invalid(api_client):
|
||||
def test_post_conversation_model_selection_new(
|
||||
api_client,
|
||||
mock_openai_stream,
|
||||
mock_uuid4,
|
||||
settings,
|
||||
):
|
||||
"""Test the user can select a different model."""
|
||||
@@ -1061,14 +1015,10 @@ def test_post_conversation_model_selection_new(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"Hello"\n'
|
||||
'0:" there"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1082,6 +1032,7 @@ def test_post_conversation_model_selection_new(
|
||||
def test_post_conversation_data_protocol_no_stream(
|
||||
api_client,
|
||||
mock_openai_no_stream,
|
||||
mock_uuid4,
|
||||
settings,
|
||||
stream_delay,
|
||||
):
|
||||
@@ -1133,9 +1084,6 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
# Wait for the content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
if stream_delay:
|
||||
assert response_content == (
|
||||
'0:"The "\n'
|
||||
@@ -1155,13 +1103,13 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
'0:" sca"\n'
|
||||
'0:"tter"\n'
|
||||
'0:"ing."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":135}}\n'
|
||||
)
|
||||
else:
|
||||
assert response_content == (
|
||||
'0:"The sky appears blue due to a phenomenon called Rayleigh scattering."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":135}}\n'
|
||||
)
|
||||
|
||||
@@ -1180,9 +1128,8 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Why the sky is blue?",
|
||||
reasoning=None,
|
||||
@@ -1193,9 +1140,8 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
parts=[TextUIPart(type="text", text="Why the sky is blue?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The sky appears blue due to a phenomenon called Rayleigh scattering.",
|
||||
reasoning=None,
|
||||
@@ -1211,7 +1157,6 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
@@ -1241,7 +1186,6 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -1269,154 +1213,5 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 135,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_conversation_async(api_client, mock_openai_stream, monkeypatch, caplog):
|
||||
"""Test posting messages to a conversation using the 'data' protocol."""
|
||||
monkeypatch.setenv("PYTHON_SERVER_MODE", "async")
|
||||
|
||||
chat_conversation = await sync_to_async(ChatConversationFactory)(owner__language="en-us")
|
||||
|
||||
url = f"/api/v1.0/chats/{chat_conversation.pk}/conversation/?protocol=data"
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"id": "yuPoOuBkKA4FnKvk",
|
||||
"role": "user",
|
||||
"parts": [{"text": "Hello", "type": "text"}],
|
||||
"content": "Hello",
|
||||
"createdAt": "2025-07-03T15:22:17.105Z",
|
||||
}
|
||||
]
|
||||
}
|
||||
await api_client.aforce_login(chat_conversation.owner)
|
||||
|
||||
caplog.clear()
|
||||
caplog.set_level(level=logging.DEBUG, logger="chat.views")
|
||||
|
||||
response = await sync_to_async(api_client.post)(url, data, format="json") # client is sync
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.get("Content-Type") == "text/event-stream"
|
||||
assert response.get("x-vercel-ai-data-stream") == "v1"
|
||||
assert response.streaming
|
||||
|
||||
assert "Using ASYNC streaming for chat conversation" in caplog.text
|
||||
|
||||
# Wait for the streaming content to be fully received => async iterator -> list
|
||||
# This fails it the streaming is not an async generator
|
||||
response_content = b"".join([content async for content in response.streaming_content]).decode(
|
||||
"utf-8"
|
||||
)
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"Hello"\n'
|
||||
'0:" there"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
assert mock_openai_stream.called
|
||||
|
||||
await chat_conversation.arefresh_from_db()
|
||||
assert chat_conversation.ui_messages == [
|
||||
{
|
||||
"content": "Hello",
|
||||
"createdAt": "2025-07-03T15:22:17.105Z",
|
||||
"id": "yuPoOuBkKA4FnKvk",
|
||||
"parts": [{"text": "Hello", "type": "text"}],
|
||||
"role": "user",
|
||||
}
|
||||
]
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="Hello there")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "Today is Friday 25/07/2025.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": ["Hello"],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
"kind": "response",
|
||||
"model_name": "test-model",
|
||||
"parts": [{"content": "Hello there", "id": None, "part_kind": "text"}],
|
||||
"provider_details": {"finish_reason": "stop"},
|
||||
"provider_name": "openai",
|
||||
"provider_response_id": "chatcmpl-1234567890",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"details": {},
|
||||
"input_audio_tokens": 0,
|
||||
"input_tokens": 0,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
+80
-169
@@ -8,13 +8,12 @@ import logging
|
||||
from io import BytesIO
|
||||
from unittest import mock
|
||||
|
||||
from django.utils import formats, timezone
|
||||
from django.utils import timezone
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import responses
|
||||
import respx
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
|
||||
from pydantic_ai.models.function import AgentInfo, DeltaToolCall, FunctionModel
|
||||
@@ -33,7 +32,6 @@ from chat.ai_sdk_types import (
|
||||
UIMessage,
|
||||
)
|
||||
from chat.factories import ChatConversationFactory
|
||||
from chat.tests.utils import replace_uuids_with_placeholder
|
||||
|
||||
# enable database transactions for tests:
|
||||
# transaction=True ensures that the data are available in the database
|
||||
@@ -151,7 +149,7 @@ def fixture_mock_summarization_agent():
|
||||
super().__init__(**kwargs)
|
||||
self._model = FunctionModel(function=summarization_model) # pylint: disable=protected-access
|
||||
|
||||
with mock.patch("chat.tools.document_summarize.SummarizationAgent", new=SummarizationAgentMock):
|
||||
with mock.patch("chat.agents.summarize.SummarizationAgent", new=SummarizationAgentMock):
|
||||
yield
|
||||
|
||||
|
||||
@@ -215,12 +213,12 @@ def fixture_mock_openai_stream():
|
||||
|
||||
@responses.activate
|
||||
@respx.mock
|
||||
@freeze_time()
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def test_post_conversation_with_document_upload( # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
api_client,
|
||||
mock_albert_api, # pylint: disable=unused-argument
|
||||
sample_pdf_content,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -274,11 +272,9 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
str_mock_uuid4 = str(mock_uuid4)
|
||||
toolcall_id = f"pyd_ai_{str_mock_uuid4.replace('-', '')}"
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'9:{"toolCallId":"XXX","toolName":"document_parsing",'
|
||||
'"args":{"documents":[{"identifier":"sample.pdf"}]}}\n'
|
||||
@@ -286,22 +282,19 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
'b:{"toolCallId":"pyd_ai_YYY","toolName":"document_search_rag"}\n'
|
||||
'9:{"toolCallId":"pyd_ai_YYY","toolName":"document_search_rag",'
|
||||
'"args":{"query":"What does the document say?"}}\n'
|
||||
'h:{"sourceType":"url","id":"<mocked_uuid>","url":"sample.pdf","title":null,'
|
||||
'"providerMetadata":{}}\n'
|
||||
'h:{"sourceType":"url","id":"XXX","url":"sample.pdf","title":null,"providerMetadata":{}}\n'
|
||||
'a:{"toolCallId":"pyd_ai_YYY","result":[{"url":"sample.pdf","content":"This '
|
||||
'is the content of the PDF.","score":0.9}]}\n'
|
||||
"0:\"From the document, I can see that it says 'Hello PDF'.\"\n"
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":100,"completionTokens":20}}\n'
|
||||
)
|
||||
).replace("XXX", str_mock_uuid4).replace("pyd_ai_YYY", toolcall_id)
|
||||
|
||||
# Check that the conversation was updated
|
||||
chat_conversation.refresh_from_db()
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id,
|
||||
id=str_mock_uuid4,
|
||||
createdAt=timezone.now(),
|
||||
content="What does the document say?",
|
||||
reasoning=None,
|
||||
@@ -311,10 +304,8 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="What does the document say?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id,
|
||||
id=str_mock_uuid4,
|
||||
createdAt=timezone.now(),
|
||||
content="From the document, I can see that it says 'Hello PDF'.",
|
||||
reasoning=None,
|
||||
@@ -326,7 +317,7 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
ToolInvocationUIPart(
|
||||
type="tool-invocation",
|
||||
toolInvocation=ToolInvocationCall(
|
||||
toolCallId=chat_conversation.messages[1].parts[0].toolInvocation.toolCallId,
|
||||
toolCallId=toolcall_id,
|
||||
toolName="document_search_rag",
|
||||
args={"query": "What does the document say?"},
|
||||
state="call",
|
||||
@@ -338,7 +329,7 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
type="source",
|
||||
source=LanguageModelV1Source(
|
||||
sourceType="url",
|
||||
id=chat_conversation.messages[1].parts[2].source.id,
|
||||
id=str_mock_uuid4,
|
||||
url="sample.pdf",
|
||||
title=None,
|
||||
providerMetadata={},
|
||||
@@ -347,66 +338,55 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
],
|
||||
)
|
||||
|
||||
timezone_now = timezone.now().isoformat().replace("+00:00", "Z")
|
||||
_formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
|
||||
assert len(chat_conversation.pydantic_messages) == 4
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages[0] == {
|
||||
"instructions": "When you receive a result from the summarization tool, you "
|
||||
"MUST return it directly to the user without any "
|
||||
"modification, paraphrasing, or additional summarization.The "
|
||||
"tool already produces optimized summaries that should be "
|
||||
"presented verbatim.You may translate the summary if "
|
||||
"required, but you MUST preserve all the information from the "
|
||||
"original summary.You may add a follow-up question after the "
|
||||
"summary if needed.",
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"content": "Today is Friday 25/07/2025.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "Use document_search_rag ONLY to retrieve specific "
|
||||
"passages from attached documents. Do NOT use it to "
|
||||
"summarize; for summaries, call the summarize tool "
|
||||
"instead.",
|
||||
"content": "If the user wants specific information from a "
|
||||
"document, invoke web_search_albert_rag with an "
|
||||
"appropriate query string.Do not ask the user for the "
|
||||
"document; rely on the tool to locate and return "
|
||||
"relevant passages.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "[Internal context] User documents are attached to this "
|
||||
"conversation. Do not request re-upload of documents; "
|
||||
"consider them already available via the internal "
|
||||
"store.",
|
||||
"content": "If the user wants a summary of document(s), invoke "
|
||||
"summarize tool without asking the user for the "
|
||||
"document itself. The tool will handle any necessary "
|
||||
"extraction and summarization based on the internal "
|
||||
"context.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": ["What does the document say?"],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[1] == {
|
||||
"finish_reason": None,
|
||||
@@ -417,14 +397,14 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
"args": '{"query": "What does the document say?"}',
|
||||
"id": None,
|
||||
"part_kind": "tool-call",
|
||||
"tool_call_id": chat_conversation.pydantic_messages[1]["parts"][0]["tool_call_id"],
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_name": "document_search_rag",
|
||||
}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
@@ -435,19 +415,9 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 8,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[2] == {
|
||||
"instructions": (
|
||||
"When you receive a result from the summarization tool, you MUST "
|
||||
"return it directly to the user without any modification, "
|
||||
"paraphrasing, or additional summarization."
|
||||
"The tool already produces optimized summaries that should "
|
||||
"be presented verbatim."
|
||||
"You may translate the summary if required, but you MUST preserve "
|
||||
"all the information from the original summary."
|
||||
"You may add a follow-up question after the summary if needed."
|
||||
),
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -460,12 +430,11 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
],
|
||||
"metadata": {"sources": ["sample.pdf"]},
|
||||
"part_kind": "tool-return",
|
||||
"timestamp": timezone_now,
|
||||
"tool_call_id": chat_conversation.pydantic_messages[2]["parts"][0]["tool_call_id"],
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_name": "document_search_rag",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[3] == {
|
||||
"finish_reason": None,
|
||||
@@ -481,7 +450,7 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
@@ -492,19 +461,19 @@ def test_post_conversation_with_document_upload( # pylint: disable=too-many-arg
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 12,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
|
||||
@responses.activate
|
||||
@respx.mock
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def test_post_conversation_with_document_upload_feature_disabled( # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
def test_post_conversation_with_document_upload_feature_disabled( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
api_client,
|
||||
caplog,
|
||||
mock_openai_stream, # pylint: disable=unused-argument
|
||||
sample_pdf_content,
|
||||
feature_flags,
|
||||
mock_uuid4,
|
||||
):
|
||||
"""
|
||||
Test POST to /api/v1/chats/{pk}/conversation/ with a PDF document while feature is disabled.
|
||||
@@ -549,14 +518,10 @@ def test_post_conversation_with_document_upload_feature_disabled( # pylint: dis
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"From the document, I can see that "\n'
|
||||
"0:\"it says 'Hello PDF'.\"\n"
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":150,"completionTokens":25}}\n'
|
||||
)
|
||||
|
||||
@@ -566,12 +531,12 @@ def test_post_conversation_with_document_upload_feature_disabled( # pylint: dis
|
||||
|
||||
@responses.activate
|
||||
@respx.mock
|
||||
@freeze_time()
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def test_post_conversation_with_document_upload_summarize( # pylint: disable=too-many-arguments,too-many-positional-arguments # noqa: PLR0913
|
||||
api_client,
|
||||
mock_albert_api, # pylint: disable=unused-argument
|
||||
sample_pdf_content,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
mock_summarization_agent, # pylint: disable=unused-argument
|
||||
):
|
||||
@@ -626,33 +591,26 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
str_mock_uuid4 = str(mock_uuid4)
|
||||
toolcall_id = f"pyd_ai_{str_mock_uuid4.replace('-', '')}"
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'9:{"toolCallId":"XXX","toolName":"document_parsing",'
|
||||
'"args":{"documents":[{"identifier":"sample.pdf"}]}}\n'
|
||||
'a:{"toolCallId":"XXX","result":{"state":"done"}}\n'
|
||||
'b:{"toolCallId":"pyd_ai_YYY","toolName":"summarize"}\n'
|
||||
'9:{"toolCallId":"pyd_ai_YYY","toolName":"summarize","args":{}}\n'
|
||||
'h:{"sourceType":"url","id":"<mocked_uuid>","url":"sample.pdf.md",'
|
||||
'"title":null,"providerMetadata":{}}\n'
|
||||
'a:{"toolCallId":"pyd_ai_YYY","result":"The '
|
||||
'document discusses various topics."}\n'
|
||||
'f:{"messageId":"XXX"}\n'
|
||||
'a:{"toolCallId":"pyd_ai_YYY","result":{"state":"done"}}\n'
|
||||
'0:"The document discusses various topics."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":317,"completionTokens":19}}\n'
|
||||
)
|
||||
'h:{"sourceType":"url","id":"XXX","url":"sample.pdf","title":null,"providerMetadata":{}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":150,"completionTokens":7}}\n'
|
||||
).replace("XXX", str_mock_uuid4).replace("pyd_ai_YYY", toolcall_id)
|
||||
|
||||
# Check that the conversation was updated
|
||||
chat_conversation.refresh_from_db()
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id,
|
||||
id=str_mock_uuid4,
|
||||
createdAt=timezone.now(),
|
||||
content="Make a summary of this document.",
|
||||
reasoning=None,
|
||||
@@ -662,10 +620,8 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="Make a summary of this document.")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id,
|
||||
id=str_mock_uuid4,
|
||||
createdAt=timezone.now(),
|
||||
content="The document discusses various topics.",
|
||||
reasoning=None,
|
||||
@@ -677,7 +633,7 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
ToolInvocationUIPart(
|
||||
type="tool-invocation",
|
||||
toolInvocation=ToolInvocationCall(
|
||||
toolCallId=chat_conversation.messages[1].parts[0].toolInvocation.toolCallId,
|
||||
toolCallId=toolcall_id,
|
||||
toolName="summarize",
|
||||
args={},
|
||||
state="call",
|
||||
@@ -689,8 +645,8 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
type="source",
|
||||
source=LanguageModelV1Source(
|
||||
sourceType="url",
|
||||
id=chat_conversation.messages[1].parts[2].source.id,
|
||||
url="sample.pdf.md", # might be fixed in the future
|
||||
id=str_mock_uuid4,
|
||||
url="sample.pdf",
|
||||
title=None,
|
||||
providerMetadata={},
|
||||
),
|
||||
@@ -698,66 +654,55 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
],
|
||||
)
|
||||
|
||||
timezone_now = timezone.now().isoformat().replace("+00:00", "Z")
|
||||
_formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
|
||||
assert len(chat_conversation.pydantic_messages) == 4
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert len(chat_conversation.pydantic_messages) == 3
|
||||
assert chat_conversation.pydantic_messages[0] == {
|
||||
"instructions": "When you receive a result from the summarization tool, you "
|
||||
"MUST return it directly to the user without any "
|
||||
"modification, paraphrasing, or additional summarization.The "
|
||||
"tool already produces optimized summaries that should be "
|
||||
"presented verbatim.You may translate the summary if "
|
||||
"required, but you MUST preserve all the information from the "
|
||||
"original summary.You may add a follow-up question after the "
|
||||
"summary if needed.",
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"content": "Today is Friday 25/07/2025.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "Use document_search_rag ONLY to retrieve specific "
|
||||
"passages from attached documents. Do NOT use it to "
|
||||
"summarize; for summaries, call the summarize tool "
|
||||
"instead.",
|
||||
"content": "If the user wants specific information from a "
|
||||
"document, invoke web_search_albert_rag with an "
|
||||
"appropriate query string.Do not ask the user for the "
|
||||
"document; rely on the tool to locate and return "
|
||||
"relevant passages.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": "[Internal context] User documents are attached to this "
|
||||
"conversation. Do not request re-upload of documents; "
|
||||
"consider them already available via the internal "
|
||||
"store.",
|
||||
"content": "If the user wants a summary of document(s), invoke "
|
||||
"summarize tool without asking the user for the "
|
||||
"document itself. The tool will handle any necessary "
|
||||
"extraction and summarization based on the internal "
|
||||
"context.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": ["Make a summary of this document."],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[1] == {
|
||||
"finish_reason": None,
|
||||
@@ -768,14 +713,14 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
"args": "{}",
|
||||
"id": None,
|
||||
"part_kind": "tool-call",
|
||||
"tool_call_id": chat_conversation.pydantic_messages[1]["parts"][0]["tool_call_id"],
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_name": "summarize",
|
||||
}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": timezone_now,
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
@@ -786,52 +731,18 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 1,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[2] == {
|
||||
"instructions": (
|
||||
"When you receive a result from the summarization tool, you MUST "
|
||||
"return it directly to the user without any modification, "
|
||||
"paraphrasing, or additional summarization."
|
||||
"The tool already produces optimized summaries that should "
|
||||
"be presented verbatim."
|
||||
"You may translate the summary if required, but you MUST preserve "
|
||||
"all the information from the original summary."
|
||||
"You may add a follow-up question after the summary if needed."
|
||||
),
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "The document discusses various topics.",
|
||||
"metadata": {"sources": ["sample.pdf.md"]},
|
||||
"content": "Final result processed.",
|
||||
"metadata": None,
|
||||
"part_kind": "tool-return",
|
||||
"timestamp": timezone_now,
|
||||
"tool_call_id": chat_conversation.pydantic_messages[2]["parts"][0]["tool_call_id"],
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_name": "summarize",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[3] == {
|
||||
"finish_reason": None,
|
||||
"kind": "response",
|
||||
"model_name": "function::agent_model",
|
||||
"parts": [
|
||||
{"content": "The document discusses various topics.", "id": None, "part_kind": "text"}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": timezone_now,
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"details": {},
|
||||
"input_audio_tokens": 0,
|
||||
"input_tokens": 50,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 6,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
-1086
File diff suppressed because it is too large
Load Diff
+32
-277
@@ -7,7 +7,6 @@ from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from rest_framework import status
|
||||
|
||||
@@ -19,7 +18,6 @@ from chat.ai_sdk_types import (
|
||||
UIMessage,
|
||||
)
|
||||
from chat.factories import ChatConversationFactory
|
||||
from chat.tests.utils import replace_uuids_with_placeholder
|
||||
|
||||
# enable database transactions for tests:
|
||||
# transaction=True ensures that the data are available in the database
|
||||
@@ -29,24 +27,12 @@ pytestmark = pytest.mark.django_db(transaction=True)
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def ai_settings(settings):
|
||||
"""
|
||||
Configure AI-related settings for tests on the provided settings object.
|
||||
|
||||
Sets test values for AI service base URL, API key, model, agent instructions, and sets
|
||||
AUTO_TITLE_AFTER_USER_MESSAGES to 999 to disable automatic title generation during tests.
|
||||
|
||||
Parameters:
|
||||
settings (object): Django settings-like object to be mutated for test configuration.
|
||||
|
||||
Returns:
|
||||
object: The same settings object with AI-related test configuration applied.
|
||||
"""
|
||||
"""Fixture to set AI service URLs for testing."""
|
||||
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
|
||||
settings.AI_API_KEY = "test-api-key"
|
||||
settings.AI_MODEL = "test-model"
|
||||
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
|
||||
|
||||
settings.AUTO_TITLE_AFTER_USER_MESSAGES = 999 # disable auto title generation
|
||||
return settings
|
||||
|
||||
|
||||
@@ -214,7 +200,7 @@ def history_conversation_fixture():
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_data_protocol_with_history(
|
||||
api_client, mock_openai_stream, history_conversation
|
||||
api_client, mock_openai_stream, mock_uuid4, history_conversation
|
||||
):
|
||||
"""Test posting messages to a conversation with history using the 'data' protocol."""
|
||||
url = f"/api/v1.0/chats/{history_conversation.pk}/conversation/?protocol=data"
|
||||
@@ -240,14 +226,10 @@ def test_post_conversation_data_protocol_with_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"Hello"\n'
|
||||
'0:" there"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -277,9 +259,8 @@ def test_post_conversation_data_protocol_with_history(
|
||||
assert len(history_conversation.messages) == 6
|
||||
|
||||
# Verify the most recent message is the new one
|
||||
assert history_conversation.messages[4].id == IsUUID(4)
|
||||
assert history_conversation.messages[4] == UIMessage(
|
||||
id=history_conversation.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -290,9 +271,8 @@ def test_post_conversation_data_protocol_with_history(
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert history_conversation.messages[5].id == IsUUID(4)
|
||||
assert history_conversation.messages[5] == UIMessage(
|
||||
id=history_conversation.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -310,7 +290,7 @@ def test_post_conversation_data_protocol_with_history(
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_text_protocol_with_history(
|
||||
api_client, mock_openai_stream, history_conversation
|
||||
api_client, mock_openai_stream, mock_uuid4, history_conversation
|
||||
):
|
||||
"""Test posting messages to a conversation with history using the 'text' protocol."""
|
||||
url = f"/api/v1.0/chats/{history_conversation.pk}/conversation/?protocol=text"
|
||||
@@ -355,9 +335,8 @@ def test_post_conversation_text_protocol_with_history(
|
||||
assert len(history_conversation.messages) == 6
|
||||
|
||||
# Verify the most recent messages are the new ones
|
||||
assert history_conversation.messages[4].id == IsUUID(4)
|
||||
assert history_conversation.messages[4] == UIMessage(
|
||||
id=history_conversation.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -368,9 +347,8 @@ def test_post_conversation_text_protocol_with_history(
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert history_conversation.messages[5].id == IsUUID(4)
|
||||
assert history_conversation.messages[5] == UIMessage(
|
||||
id=history_conversation.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -385,7 +363,7 @@ def test_post_conversation_text_protocol_with_history(
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_with_image_with_history(
|
||||
api_client, mock_openai_stream_image, history_conversation
|
||||
api_client, mock_openai_stream_image, mock_uuid4, history_conversation
|
||||
):
|
||||
"""
|
||||
Ensure an image URL is correctly forwarded to the AI service with a conversation with history.
|
||||
@@ -425,14 +403,10 @@ def test_post_conversation_with_image_with_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"I see a cat"\n'
|
||||
'0:" in the picture."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -478,9 +452,8 @@ def test_post_conversation_with_image_with_history(
|
||||
assert len(history_conversation.messages) == 6
|
||||
|
||||
# Verify the most recent message has the image attachment
|
||||
assert history_conversation.messages[4].id == IsUUID(4)
|
||||
assert history_conversation.messages[4] == UIMessage(
|
||||
id=history_conversation.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello, what do you see on this picture?",
|
||||
reasoning=None,
|
||||
@@ -501,9 +474,8 @@ def test_post_conversation_with_image_with_history(
|
||||
parts=[TextUIPart(type="text", text="Hello, what do you see on this picture?")],
|
||||
)
|
||||
|
||||
assert history_conversation.messages[5].id == IsUUID(4)
|
||||
assert history_conversation.messages[5] == UIMessage(
|
||||
id=history_conversation.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I see a cat in the picture.",
|
||||
reasoning=None,
|
||||
@@ -518,7 +490,7 @@ def test_post_conversation_with_image_with_history(
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_tool_call_with_history(
|
||||
api_client, mock_openai_stream_tool, settings, history_conversation
|
||||
api_client, mock_openai_stream_tool, mock_uuid4, settings, history_conversation
|
||||
):
|
||||
"""
|
||||
Ensure tool calls are correctly forwarded and streamed back with a conversation with history.
|
||||
@@ -549,10 +521,6 @@ def test_post_conversation_tool_call_with_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'b:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","toolName":'
|
||||
'"get_current_weather"}\n'
|
||||
@@ -561,7 +529,7 @@ def test_post_conversation_tool_call_with_history(
|
||||
'a:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","result":{"location":'
|
||||
'"Paris","temperature":22,"unit":"celsius"}}\n'
|
||||
'0:"The current weather in Paris is nice"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -593,9 +561,8 @@ def test_post_conversation_tool_call_with_history(
|
||||
assert len(history_conversation.messages) == 6
|
||||
|
||||
# Verify the most recent message is the new one with tool invocation
|
||||
assert history_conversation.messages[4].id == IsUUID(4)
|
||||
assert history_conversation.messages[4] == UIMessage(
|
||||
id=history_conversation.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -606,9 +573,8 @@ def test_post_conversation_tool_call_with_history(
|
||||
parts=[TextUIPart(type="text", text="Weather in Paris?")],
|
||||
)
|
||||
|
||||
assert history_conversation.messages[5].id == IsUUID(4)
|
||||
assert history_conversation.messages[5] == UIMessage(
|
||||
id=history_conversation.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The current weather in Paris is nice",
|
||||
reasoning=None,
|
||||
@@ -640,7 +606,7 @@ def test_post_conversation_tool_call_with_history(
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_tool_call_fails_with_history(
|
||||
api_client, mock_openai_stream_tool, settings, history_conversation
|
||||
api_client, mock_openai_stream_tool, mock_uuid4, settings, history_conversation
|
||||
):
|
||||
"""
|
||||
Ensure tool calls are correctly forwarded and streamed back when failing with a
|
||||
@@ -672,10 +638,6 @@ def test_post_conversation_tool_call_fails_with_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'b:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","toolName":'
|
||||
'"get_current_weather"}\n'
|
||||
@@ -684,7 +646,7 @@ def test_post_conversation_tool_call_fails_with_history(
|
||||
'a:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","result":"Unknown tool '
|
||||
"name: 'get_current_weather'. No tools available.\"}\n"
|
||||
'0:"I cannot give you an answer to that."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -716,9 +678,8 @@ def test_post_conversation_tool_call_fails_with_history(
|
||||
assert len(history_conversation.messages) == 6
|
||||
|
||||
# Verify the most recent message is the new one with tool invocation
|
||||
assert history_conversation.messages[4].id == IsUUID(4)
|
||||
assert history_conversation.messages[4] == UIMessage(
|
||||
id=history_conversation.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -729,9 +690,8 @@ def test_post_conversation_tool_call_fails_with_history(
|
||||
parts=[TextUIPart(type="text", text="Weather in Paris?")],
|
||||
)
|
||||
|
||||
assert history_conversation.messages[5].id == IsUUID(4)
|
||||
assert history_conversation.messages[5] == UIMessage(
|
||||
id=history_conversation.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I cannot give you an answer to that.",
|
||||
reasoning=None,
|
||||
@@ -1187,7 +1147,7 @@ def history_conversation_with_tool_fixture():
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_with_existing_image_history(
|
||||
api_client, mock_openai_stream, history_conversation_with_image
|
||||
api_client, mock_openai_stream, mock_uuid4, history_conversation_with_image
|
||||
):
|
||||
"""Test posting a message to a conversation that already has images in its history."""
|
||||
url = f"/api/v1.0/chats/{history_conversation_with_image.pk}/conversation/?protocol=data"
|
||||
@@ -1213,14 +1173,10 @@ def test_post_conversation_with_existing_image_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"Hello"\n'
|
||||
'0:" there"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1251,9 +1207,8 @@ def test_post_conversation_with_existing_image_history(
|
||||
assert len(history_conversation_with_image.messages) == 6
|
||||
|
||||
# Verify the most recent messages are the new ones
|
||||
assert history_conversation_with_image.messages[4].id == IsUUID(4)
|
||||
assert history_conversation_with_image.messages[4] == UIMessage(
|
||||
id=history_conversation_with_image.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="What was in that image again?",
|
||||
reasoning=None,
|
||||
@@ -1264,9 +1219,8 @@ def test_post_conversation_with_existing_image_history(
|
||||
parts=[TextUIPart(type="text", text="What was in that image again?")],
|
||||
)
|
||||
|
||||
assert history_conversation_with_image.messages[5].id == IsUUID(4)
|
||||
assert history_conversation_with_image.messages[5] == UIMessage(
|
||||
id=history_conversation_with_image.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -1284,7 +1238,7 @@ def test_post_conversation_with_existing_image_history(
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_with_existing_tool_history(
|
||||
api_client, mock_openai_stream_tool, settings, history_conversation_with_tool
|
||||
api_client, mock_openai_stream_tool, mock_uuid4, settings, history_conversation_with_tool
|
||||
):
|
||||
"""Test posting a message to a conversation that already has tool calls in its history."""
|
||||
settings.AI_AGENT_TOOLS = ["get_current_weather"]
|
||||
@@ -1312,10 +1266,6 @@ def test_post_conversation_with_existing_tool_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'b:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","toolName":'
|
||||
'"get_current_weather"}\n'
|
||||
@@ -1324,7 +1274,7 @@ def test_post_conversation_with_existing_tool_history(
|
||||
'a:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","result":{"location":'
|
||||
'"Paris","temperature":22,"unit":"celsius"}}\n'
|
||||
'0:"The current weather in Paris is nice"\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1344,9 +1294,8 @@ def test_post_conversation_with_existing_tool_history(
|
||||
assert len(history_conversation_with_tool.messages) == 6
|
||||
|
||||
# Verify the most recent message is the new one with tool invocation
|
||||
assert history_conversation_with_tool.messages[4].id == IsUUID(4)
|
||||
assert history_conversation_with_tool.messages[4] == UIMessage(
|
||||
id=history_conversation_with_tool.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="How about Paris weather?",
|
||||
reasoning=None,
|
||||
@@ -1357,9 +1306,8 @@ def test_post_conversation_with_existing_tool_history(
|
||||
parts=[TextUIPart(type="text", text="How about Paris weather?")],
|
||||
)
|
||||
|
||||
assert history_conversation_with_tool.messages[5].id == IsUUID(4)
|
||||
assert history_conversation_with_tool.messages[5] == UIMessage(
|
||||
id=history_conversation_with_tool.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The current weather in Paris is nice",
|
||||
reasoning=None,
|
||||
@@ -1385,8 +1333,6 @@ def test_post_conversation_with_existing_tool_history(
|
||||
# The pydantic_messages should include both the original tool calls and the new ones
|
||||
assert len(history_conversation_with_tool.pydantic_messages) == 12 # Original 8 + 4 new ones
|
||||
|
||||
_run_id = history_conversation_with_tool.pydantic_messages[8]["run_id"]
|
||||
|
||||
# Verify the new tool call request is included
|
||||
assert history_conversation_with_tool.pydantic_messages[8] == {
|
||||
"instructions": None,
|
||||
@@ -1398,7 +1344,6 @@ def test_post_conversation_with_existing_tool_history(
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
assert history_conversation_with_tool.pydantic_messages[9] == {
|
||||
@@ -1428,7 +1373,6 @@ def test_post_conversation_with_existing_tool_history(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
assert history_conversation_with_tool.pydantic_messages[10] == {
|
||||
@@ -1444,7 +1388,6 @@ def test_post_conversation_with_existing_tool_history(
|
||||
"tool_name": "get_current_weather",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
assert history_conversation_with_tool.pydantic_messages[11] == {
|
||||
@@ -1468,14 +1411,13 @@ def test_post_conversation_with_existing_tool_history(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_add_image_to_conversation_with_tool_history(
|
||||
api_client, mock_openai_stream_image, history_conversation_with_tool
|
||||
api_client, mock_openai_stream_image, mock_uuid4, history_conversation_with_tool
|
||||
):
|
||||
"""Test adding an image to a conversation that already has tool calls in its history."""
|
||||
url = f"/api/v1.0/chats/{history_conversation_with_tool.pk}/conversation/?protocol=data"
|
||||
@@ -1513,14 +1455,10 @@ def test_post_conversation_add_image_to_conversation_with_tool_history(
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"I see a cat"\n'
|
||||
'0:" in the picture."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1546,9 +1484,8 @@ def test_post_conversation_add_image_to_conversation_with_tool_history(
|
||||
assert len(history_conversation_with_tool.messages) == 6
|
||||
|
||||
# Verify the most recent message has the image attachment
|
||||
assert history_conversation_with_tool.messages[4].id == IsUUID(4)
|
||||
assert history_conversation_with_tool.messages[4] == UIMessage(
|
||||
id=history_conversation_with_tool.messages[4].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="How's the weather in this image?",
|
||||
reasoning=None,
|
||||
@@ -1569,9 +1506,8 @@ def test_post_conversation_add_image_to_conversation_with_tool_history(
|
||||
parts=[TextUIPart(type="text", text="How's the weather in this image?")],
|
||||
)
|
||||
|
||||
assert history_conversation_with_tool.messages[5].id == IsUUID(4)
|
||||
assert history_conversation_with_tool.messages[5] == UIMessage(
|
||||
id=history_conversation_with_tool.messages[5].id,
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I see a cat in the picture.",
|
||||
reasoning=None,
|
||||
@@ -1581,184 +1517,3 @@ def test_post_conversation_add_image_to_conversation_with_tool_history(
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="I see a cat in the picture.")],
|
||||
)
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_triggers_automatic_title_generation(
|
||||
api_client, mock_openai_stream_with_title_generation, settings, history_conversation
|
||||
):
|
||||
"""
|
||||
Test that posting the 3rd user message triggers automatic title generation.
|
||||
|
||||
The history_conversation fixture has 2 user messages. Posting a 3rd message
|
||||
should trigger title generation via the SummarizationAgent.
|
||||
"""
|
||||
# Configure the title generation threshold
|
||||
settings.AUTO_TITLE_AFTER_USER_MESSAGES = 3
|
||||
|
||||
url = f"/api/v1.0/chats/{history_conversation.pk}/conversation/?protocol=data"
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"id": "third-user-msg",
|
||||
"role": "user",
|
||||
"parts": [{"text": "Can you explain backpropagation?", "type": "text"}],
|
||||
"content": "Can you explain backpropagation?",
|
||||
"createdAt": "2025-07-25T10:36:00.000Z",
|
||||
}
|
||||
]
|
||||
}
|
||||
api_client.force_login(history_conversation.owner)
|
||||
|
||||
history_conversation.title = "initial title"
|
||||
history_conversation.save()
|
||||
|
||||
assert not history_conversation.title_set_by_user_at
|
||||
|
||||
response = api_client.post(url, data, format="json")
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.get("Content-Type") == "text/event-stream"
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Verify the conversation_metadata event is in the stream
|
||||
|
||||
assert '"type": "conversation_metadata"' in response_content
|
||||
|
||||
# Refresh and verify title was updated
|
||||
history_conversation.refresh_from_db()
|
||||
|
||||
assert history_conversation.title == "GENERATED TITLE"
|
||||
# title_set_by_user_at should remain None since it was auto-generated
|
||||
assert not history_conversation.title_set_by_user_at
|
||||
|
||||
assert mock_openai_stream_with_title_generation.called
|
||||
|
||||
assert mock_openai_stream_with_title_generation.call_count == 2
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_does_not_regenerate_title_when_user_set(
|
||||
api_client, mock_openai_stream_with_title_generation, settings, history_conversation
|
||||
):
|
||||
"""
|
||||
Test that title is NOT regenerated if the user has manually set a title.
|
||||
"""
|
||||
settings.AUTO_TITLE_AFTER_USER_MESSAGES = 3
|
||||
|
||||
# Simulate user having set a custom title
|
||||
history_conversation.title = "My Custom Title"
|
||||
history_conversation.title_set_by_user_at = timezone.now()
|
||||
history_conversation.save()
|
||||
|
||||
url = f"/api/v1.0/chats/{history_conversation.pk}/conversation/?protocol=data"
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"id": "third-user-msg",
|
||||
"role": "user",
|
||||
"parts": [{"text": "Can you explain backpropagation?", "type": "text"}],
|
||||
"content": "Can you explain backpropagation?",
|
||||
"createdAt": "2025-07-25T10:36:00.000Z",
|
||||
}
|
||||
]
|
||||
}
|
||||
api_client.force_login(history_conversation.owner)
|
||||
|
||||
response = api_client.post(url, data, format="json")
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
|
||||
# Consume the stream
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# conversation_metadata should NOT be in the stream since title wasn't generated
|
||||
assert "conversation_metadata" not in response_content
|
||||
|
||||
# Refresh and verify title was NOT changed
|
||||
history_conversation.refresh_from_db()
|
||||
|
||||
assert history_conversation.title == "My Custom Title"
|
||||
assert history_conversation.title_set_by_user_at is not None
|
||||
|
||||
assert mock_openai_stream_with_title_generation.called
|
||||
|
||||
assert mock_openai_stream_with_title_generation.call_count == 1
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
def test_post_conversation_does_not_generate_title_before_threshold(
|
||||
api_client, mock_openai_stream_with_title_generation, settings
|
||||
):
|
||||
"""
|
||||
Test that title is NOT generated before reaching the message threshold.
|
||||
"""
|
||||
settings.AUTO_TITLE_AFTER_USER_MESSAGES = 3
|
||||
|
||||
# Create a conversation with only 1 user message
|
||||
history_timestamp = timezone.now().replace(year=2025, month=6, day=15, hour=10, minute=30)
|
||||
conversation = ChatConversationFactory(title="initial title")
|
||||
|
||||
conversation.messages = [
|
||||
UIMessage(
|
||||
id="prev-user-msg-1",
|
||||
createdAt=history_timestamp,
|
||||
content="Hello!",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="Hello!")],
|
||||
),
|
||||
UIMessage(
|
||||
id="prev-assistant-msg-1",
|
||||
createdAt=history_timestamp.replace(minute=31),
|
||||
content="Hi there! How can I help you?",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[TextUIPart(type="text", text="Hi there! How can I help you?")],
|
||||
),
|
||||
]
|
||||
conversation.save()
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk}/conversation/?protocol=data"
|
||||
data = {
|
||||
"messages": [
|
||||
{
|
||||
"id": "second-user-msg",
|
||||
"role": "user",
|
||||
"parts": [{"text": "What's machine learning?", "type": "text"}],
|
||||
"content": "What's machine learning?",
|
||||
"createdAt": "2025-07-25T10:36:00.000Z",
|
||||
}
|
||||
]
|
||||
}
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
response = api_client.post(url, data, format="json")
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
|
||||
# Consume the stream
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# conversation_metadata should NOT be in the stream (only 2 user messages)
|
||||
assert "conversation_metadata" not in response_content
|
||||
|
||||
# Refresh and verify title was not updated
|
||||
conversation.refresh_from_db()
|
||||
|
||||
assert conversation.title == "initial title"
|
||||
assert not conversation.title_set_by_user_at
|
||||
|
||||
assert mock_openai_stream_with_title_generation.call_count == 1
|
||||
@@ -1,837 +0,0 @@
|
||||
"""Unit tests for chat conversation actions with image URL."""
|
||||
|
||||
import uuid
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from pydantic_ai import ModelRequest, RequestUsage
|
||||
from pydantic_ai.messages import (
|
||||
ImageUrl,
|
||||
ModelMessage,
|
||||
ModelResponse,
|
||||
SystemPromptPart,
|
||||
TextPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models.function import AgentInfo, FunctionModel
|
||||
from rest_framework import status
|
||||
|
||||
from chat.ai_sdk_types import (
|
||||
Attachment,
|
||||
TextUIPart,
|
||||
UIMessage,
|
||||
)
|
||||
from chat.factories import ChatConversationFactory
|
||||
from chat.tests.utils import replace_uuids_with_placeholder
|
||||
|
||||
# enable database transactions for tests:
|
||||
# transaction=True ensures that the data are available in the database
|
||||
# in other threads
|
||||
pytestmark = pytest.mark.django_db(transaction=True)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def ai_settings(settings):
|
||||
"""Fixture to set AI service URLs for testing."""
|
||||
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
|
||||
settings.AI_API_KEY = "test-api-key"
|
||||
settings.AI_MODEL = "test-model"
|
||||
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
|
||||
return settings
|
||||
|
||||
|
||||
@pytest.fixture(name="sample_image_content")
|
||||
def fixture_sample_image_content():
|
||||
"""Create a dummy image content as bytes."""
|
||||
# This is a simple, valid 1x1 PNG image content.
|
||||
return (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01"
|
||||
b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\x0cIDATx\x9c\x63\x00"
|
||||
b"\x01\x00\x00\x05\x00\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
|
||||
|
||||
@freeze_time("2025-10-18T20:48:20.286204Z")
|
||||
def test_post_conversation_with_local_image_url(
|
||||
api_client,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
Test POST to /api/v1/chats/{pk}/conversation/ with an image URL.
|
||||
"""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
image_url = f"/media-key/{chat_conversation.pk}/sample.png"
|
||||
|
||||
message = UIMessage(
|
||||
id="1",
|
||||
role="user",
|
||||
content="What is in this image?",
|
||||
parts=[
|
||||
TextUIPart(
|
||||
text="What is in this image?",
|
||||
type="text",
|
||||
),
|
||||
],
|
||||
experimental_attachments=[
|
||||
Attachment(
|
||||
name="sample.png",
|
||||
contentType="image/png",
|
||||
url=image_url,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
presigned_url = messages[0].parts[3].content[1].url
|
||||
assert presigned_url.startswith("http://localhost:9000/conversations-media-storage/")
|
||||
assert presigned_url.find("X-Amz-Signature=") != -1
|
||||
assert presigned_url.find("X-Amz-Date=") != -1
|
||||
assert presigned_url.find("X-Amz-Expires=") != -1
|
||||
|
||||
assert messages == [
|
||||
ModelRequest(
|
||||
parts=[
|
||||
SystemPromptPart(
|
||||
content="You are a helpful test assistant :)", timestamp=timezone.now()
|
||||
),
|
||||
SystemPromptPart(
|
||||
content="Today is Saturday 18/10/2025.", timestamp=timezone.now()
|
||||
),
|
||||
SystemPromptPart(content="Answer in english.", timestamp=timezone.now()),
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"What is in this image?",
|
||||
ImageUrl(
|
||||
url=presigned_url, # presigned URL for this conversation
|
||||
media_type="image/png",
|
||||
identifier="sample.png",
|
||||
),
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
],
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "This is an image of a single pixel."
|
||||
|
||||
# Use the fixture with FunctionModel
|
||||
with mock_ai_agent_service(FunctionModel(stream_function=agent_model)):
|
||||
response = api_client.post(
|
||||
f"/api/v1.0/chats/{chat_conversation.pk}/conversation/",
|
||||
data={"messages": [message.model_dump(mode="json")]},
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.get("Content-Type") == "text/event-stream"
|
||||
assert response.get("x-vercel-ai-data-stream") == "v1"
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"This is an image of a single pixel."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":50,"completionTokens":9}}\n'
|
||||
)
|
||||
|
||||
# Check that the conversation was updated
|
||||
chat_conversation.refresh_from_db()
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
experimental_attachments=[
|
||||
Attachment(name="sample.png", contentType="image/png", url=image_url)
|
||||
],
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="This is an image of a single pixel.",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="This is an image of a single pixel."),
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": "Today is Saturday 18/10/2025.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
"What is in this image?",
|
||||
{
|
||||
"force_download": False,
|
||||
"identifier": "sample.png",
|
||||
"kind": "image-url",
|
||||
"media_type": "image/png",
|
||||
"url": image_url,
|
||||
"vendor_metadata": None,
|
||||
},
|
||||
],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
"kind": "response",
|
||||
"model_name": "function::agent_model",
|
||||
"parts": [
|
||||
{"content": "This is an image of a single pixel.", "id": None, "part_kind": "text"}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"details": {},
|
||||
"input_audio_tokens": 0,
|
||||
"input_tokens": 50,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 9,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@freeze_time()
|
||||
def test_post_conversation_with_local_image_wrong_url(
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
Test POST to /api/v1/chats/{pk}/conversation/ with a tampered URL.
|
||||
"""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
image_url = f"/media-key/{uuid.uuid4()}/sample.png"
|
||||
|
||||
message = UIMessage(
|
||||
id="1",
|
||||
role="user",
|
||||
content="What is in this image?",
|
||||
parts=[
|
||||
TextUIPart(
|
||||
text="What is in this image?",
|
||||
type="text",
|
||||
),
|
||||
],
|
||||
experimental_attachments=[
|
||||
Attachment(
|
||||
name="sample.png",
|
||||
contentType="image/png",
|
||||
url=image_url,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
assert messages == [
|
||||
ModelRequest(
|
||||
parts=[
|
||||
SystemPromptPart(
|
||||
content="You are a helpful test assistant :)", timestamp=timezone.now()
|
||||
),
|
||||
SystemPromptPart(content=today_promt_date, timestamp=timezone.now()),
|
||||
SystemPromptPart(content="Answer in english.", timestamp=timezone.now()),
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"What is in this image?",
|
||||
ImageUrl(
|
||||
url=image_url, # not presigned URL for this conversation
|
||||
media_type="image/png",
|
||||
identifier="sample.png",
|
||||
),
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
],
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "cannot read image." # IRL a 400 error would be raised by the LLM
|
||||
|
||||
# Use the fixture with FunctionModel
|
||||
with mock_ai_agent_service(FunctionModel(stream_function=agent_model)):
|
||||
response = api_client.post(
|
||||
f"/api/v1.0/chats/{chat_conversation.pk}/conversation/",
|
||||
data={"messages": [message.model_dump(mode="json")]},
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.get("Content-Type") == "text/event-stream"
|
||||
assert response.get("x-vercel-ai-data-stream") == "v1"
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"cannot read image."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":50,"completionTokens":4}}\n'
|
||||
)
|
||||
|
||||
# We don't check conversation messages here because the LLM would
|
||||
# normally raise an error when trying to access the image.
|
||||
|
||||
|
||||
@freeze_time()
|
||||
def test_post_conversation_with_remote_image_url(
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
Test POST to /api/v1/chats/{pk}/conversation/ with a remote URL.
|
||||
"""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
image_url = "https://example.com/sample.png"
|
||||
|
||||
message = UIMessage(
|
||||
id="1",
|
||||
role="user",
|
||||
content="What is in this image?",
|
||||
parts=[
|
||||
TextUIPart(
|
||||
text="What is in this image?",
|
||||
type="text",
|
||||
),
|
||||
],
|
||||
experimental_attachments=[
|
||||
Attachment(
|
||||
name="sample.png",
|
||||
contentType="image/png",
|
||||
url=image_url,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
assert messages == [
|
||||
ModelRequest(
|
||||
parts=[
|
||||
SystemPromptPart(
|
||||
content="You are a helpful test assistant :)", timestamp=timezone.now()
|
||||
),
|
||||
SystemPromptPart(content=today_promt_date, timestamp=timezone.now()),
|
||||
SystemPromptPart(content="Answer in english.", timestamp=timezone.now()),
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"What is in this image?",
|
||||
ImageUrl(
|
||||
url=image_url, # remote URL
|
||||
media_type="image/png",
|
||||
identifier="sample.png",
|
||||
),
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
],
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "This is an image of a single pixel."
|
||||
|
||||
# Use the fixture with FunctionModel
|
||||
with mock_ai_agent_service(FunctionModel(stream_function=agent_model)):
|
||||
response = api_client.post(
|
||||
f"/api/v1.0/chats/{chat_conversation.pk}/conversation/",
|
||||
data={"messages": [message.model_dump(mode="json")]},
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.get("Content-Type") == "text/event-stream"
|
||||
assert response.get("x-vercel-ai-data-stream") == "v1"
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"This is an image of a single pixel."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":50,"completionTokens":9}}\n'
|
||||
)
|
||||
|
||||
# Check that the conversation was updated
|
||||
chat_conversation.refresh_from_db()
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
experimental_attachments=[
|
||||
Attachment(name="sample.png", contentType="image/png", url=image_url)
|
||||
],
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="This is an image of a single pixel.",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="This is an image of a single pixel."),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@freeze_time("2025-10-18T20:48:20.286204Z")
|
||||
def test_post_conversation_with_local_image_url_in_history(
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
Test POST to /api/v1/chats/{pk}/conversation/ with an image URL.
|
||||
"""
|
||||
chat_conversation_pk = "0be55da5-8eb7-4dad-aa0f-fea454bd5809"
|
||||
image_url = f"/media-key/{chat_conversation_pk}/sample.png"
|
||||
chat_conversation = ChatConversationFactory(
|
||||
pk=chat_conversation_pk,
|
||||
owner__language="en-us",
|
||||
messages=[
|
||||
UIMessage(
|
||||
id=str(uuid.uuid4()),
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
experimental_attachments=[
|
||||
Attachment(name="sample.png", contentType="image/png", url=image_url)
|
||||
],
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
),
|
||||
UIMessage(
|
||||
id=str(uuid.uuid4()),
|
||||
createdAt=timezone.now(),
|
||||
content="This is an image of a single pixel.",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="This is an image of a single pixel."),
|
||||
],
|
||||
),
|
||||
],
|
||||
pydantic_messages=[
|
||||
{
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
"What is in this image?",
|
||||
{
|
||||
"force_download": False,
|
||||
"identifier": "sample.png",
|
||||
"kind": "image-url",
|
||||
"media_type": "image/png",
|
||||
"url": image_url,
|
||||
"vendor_metadata": None,
|
||||
},
|
||||
],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
"kind": "response",
|
||||
"model_name": "function::agent_model",
|
||||
"parts": [
|
||||
{
|
||||
"content": "This is an image of a single pixel.",
|
||||
"id": None,
|
||||
"part_kind": "text",
|
||||
}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"details": {},
|
||||
"input_audio_tokens": 0,
|
||||
"input_tokens": 50,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 9,
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
image_url = f"/media-key/{chat_conversation.pk}/sample.png"
|
||||
|
||||
message = UIMessage(
|
||||
id="3",
|
||||
role="user",
|
||||
content="Give more details about this image.",
|
||||
parts=[
|
||||
TextUIPart(
|
||||
text="Give more details about this image.",
|
||||
type="text",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
presigned_url = messages[0].parts[3].content[1].url
|
||||
assert presigned_url.startswith("http://localhost:9000/conversations-media-storage/")
|
||||
assert presigned_url.find("X-Amz-Signature=") != -1
|
||||
assert presigned_url.find("X-Amz-Date=") != -1
|
||||
assert presigned_url.find("X-Amz-Expires=") != -1
|
||||
|
||||
assert messages == [
|
||||
ModelRequest(
|
||||
parts=[
|
||||
SystemPromptPart(
|
||||
content="You are a helpful test assistant :)",
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
SystemPromptPart(
|
||||
content=today_promt_date,
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
SystemPromptPart(
|
||||
content="Answer in english.",
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"What is in this image?",
|
||||
ImageUrl(
|
||||
url=presigned_url, # presigned URL in history
|
||||
media_type="image/png",
|
||||
identifier="sample.png",
|
||||
),
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
),
|
||||
ModelResponse(
|
||||
parts=[TextPart(content="This is an image of a single pixel.")],
|
||||
usage=RequestUsage(input_tokens=50, output_tokens=9),
|
||||
model_name="function::agent_model",
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"Give more details about this image.",
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
)
|
||||
],
|
||||
run_id=messages[2].run_id,
|
||||
),
|
||||
]
|
||||
yield "This is an image of square, very small and nice."
|
||||
|
||||
# Use the fixture with FunctionModel
|
||||
with mock_ai_agent_service(FunctionModel(stream_function=agent_model)):
|
||||
response = api_client.post(
|
||||
f"/api/v1.0/chats/{chat_conversation.pk}/conversation/",
|
||||
data={"messages": [message.model_dump(mode="json")]},
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert response.status_code == status.HTTP_200_OK
|
||||
assert response.get("Content-Type") == "text/event-stream"
|
||||
assert response.get("x-vercel-ai-data-stream") == "v1"
|
||||
assert response.streaming
|
||||
|
||||
# Wait for the streaming content to be fully received
|
||||
response_content = b"".join(response.streaming_content).decode("utf-8")
|
||||
|
||||
# Replace UUIDs with placeholders for assertion
|
||||
response_content = replace_uuids_with_placeholder(response_content)
|
||||
|
||||
assert response_content == (
|
||||
'0:"This is an image of square, very small and nice."\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":50,"completionTokens":11}}\n'
|
||||
)
|
||||
|
||||
# Check that the conversation was updated
|
||||
chat_conversation.refresh_from_db()
|
||||
assert len(chat_conversation.messages) == 2 + 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=chat_conversation.messages[0].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
experimental_attachments=[
|
||||
Attachment(name="sample.png", contentType="image/png", url=image_url)
|
||||
],
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=chat_conversation.messages[1].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="This is an image of a single pixel.",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="This is an image of a single pixel."),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[2].id == IsUUID(4)
|
||||
assert chat_conversation.messages[2] == UIMessage(
|
||||
id=chat_conversation.messages[2].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="Give more details about this image.",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="user",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="Give more details about this image."),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[3].id == IsUUID(4)
|
||||
assert chat_conversation.messages[3] == UIMessage(
|
||||
id=chat_conversation.messages[3].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="This is an image of square, very small and nice.",
|
||||
reasoning=None,
|
||||
experimental_attachments=None,
|
||||
role="assistant",
|
||||
annotations=None,
|
||||
toolInvocations=None,
|
||||
parts=[
|
||||
TextUIPart(type="text", text="This is an image of square, very small and nice."),
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[2]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
"What is in this image?",
|
||||
{
|
||||
"force_download": False,
|
||||
"identifier": "sample.png",
|
||||
"kind": "image-url",
|
||||
"media_type": "image/png",
|
||||
"url": image_url,
|
||||
"vendor_metadata": None,
|
||||
},
|
||||
],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
"kind": "response",
|
||||
"model_name": "function::agent_model",
|
||||
"parts": [
|
||||
{"content": "This is an image of a single pixel.", "id": None, "part_kind": "text"}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"details": {},
|
||||
"input_audio_tokens": 0,
|
||||
"input_tokens": 50,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 9,
|
||||
},
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": ["Give more details about this image."],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
"kind": "response",
|
||||
"model_name": "function::agent_model",
|
||||
"parts": [
|
||||
{
|
||||
"content": "This is an image of square, very small and nice.",
|
||||
"id": None,
|
||||
"part_kind": "text",
|
||||
}
|
||||
],
|
||||
"provider_details": None,
|
||||
"provider_name": None,
|
||||
"provider_response_id": None,
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
"usage": {
|
||||
"cache_audio_read_tokens": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"details": {},
|
||||
"input_audio_tokens": 0,
|
||||
"input_tokens": 50,
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 11,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
@@ -28,7 +28,6 @@ def test_create_conversation(api_client):
|
||||
conversation = ChatConversation.objects.get(id=response.data["id"])
|
||||
assert conversation.owner == user
|
||||
assert conversation.title == "New Conversation"
|
||||
assert not conversation.title_set_by_user_at
|
||||
|
||||
|
||||
def test_create_conversation_other_owner(api_client):
|
||||
|
||||
@@ -1,123 +0,0 @@
|
||||
"""Test the media authentication endpoint for chat conversations."""
|
||||
|
||||
from io import BytesIO
|
||||
from urllib.parse import urlparse
|
||||
from uuid import uuid4
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from freezegun import freeze_time
|
||||
|
||||
from core.factories import UserFactory
|
||||
|
||||
from chat.factories import ChatConversationAttachmentFactory
|
||||
from chat.models import ChatConversation
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
def test_api_media_auth_unkown_document(api_client):
|
||||
"""
|
||||
Trying to download a media related to a conversation that does not exist
|
||||
should not have the side effect to create it (no regression test).
|
||||
"""
|
||||
ChatConversation.objects.all().delete()
|
||||
|
||||
original_url = f"http://localhost/media/{uuid4()!s}/attachments/{uuid4()!s}.jpg"
|
||||
|
||||
response = api_client.get("/api/v1.0/chats/media-auth/", HTTP_X_ORIGINAL_URL=original_url)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert ChatConversation.objects.exists() is False
|
||||
|
||||
|
||||
def test_api_media_auth_anonymous(api_client):
|
||||
"""
|
||||
Users who are not owners of a conversation should not be able to retrieve
|
||||
attachments linked to it.
|
||||
"""
|
||||
attachment = ChatConversationAttachmentFactory()
|
||||
|
||||
original_url = f"http://localhost/media/{attachment.key:s}"
|
||||
response = api_client.get("/api/v1.0/chats/media-auth/", HTTP_X_ORIGINAL_URL=original_url)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert "Authorization" not in response
|
||||
|
||||
|
||||
def test_api_media_auth_owner(api_client):
|
||||
"""
|
||||
Owners of a conversation should be able to retrieve attachments linked to it.
|
||||
"""
|
||||
attachment = ChatConversationAttachmentFactory()
|
||||
|
||||
default_storage.connection.meta.client.put_object(
|
||||
Bucket=default_storage.bucket_name,
|
||||
Key=attachment.key,
|
||||
Body=BytesIO(b"my prose"),
|
||||
ContentType="text/plain",
|
||||
)
|
||||
|
||||
original_url = f"http://localhost/media/{attachment.key:s}"
|
||||
now = timezone.now()
|
||||
with freeze_time(now):
|
||||
api_client.force_login(attachment.conversation.owner)
|
||||
response = api_client.get("/api/v1.0/chats/media-auth/", HTTP_X_ORIGINAL_URL=original_url)
|
||||
|
||||
assert response.status_code == 200
|
||||
|
||||
authorization = response["Authorization"]
|
||||
assert "AWS4-HMAC-SHA256 Credential=" in authorization
|
||||
assert "SignedHeaders=host;x-amz-content-sha256;x-amz-date, Signature=" in authorization
|
||||
assert response["X-Amz-Date"] == now.strftime("%Y%m%dT%H%M%SZ")
|
||||
|
||||
s3_url = urlparse(settings.AWS_S3_ENDPOINT_URL)
|
||||
file_url = f"{settings.AWS_S3_ENDPOINT_URL:s}/conversations-media-storage/{attachment.key:s}"
|
||||
response = requests.get(
|
||||
file_url,
|
||||
headers={
|
||||
"authorization": authorization,
|
||||
"x-amz-date": response["x-amz-date"],
|
||||
"x-amz-content-sha256": response["x-amz-content-sha256"],
|
||||
"Host": f"{s3_url.hostname:s}:{s3_url.port:d}",
|
||||
},
|
||||
timeout=1,
|
||||
)
|
||||
assert response.content.decode("utf-8") == "my prose"
|
||||
|
||||
|
||||
def test_api_media_auth_not_owner(api_client):
|
||||
"""
|
||||
Users who are not owners of a conversation should not be able to retrieve
|
||||
attachments linked to it.
|
||||
"""
|
||||
attachment = ChatConversationAttachmentFactory()
|
||||
user = UserFactory()
|
||||
|
||||
api_client.force_login(user)
|
||||
original_url = f"http://localhost/media/{attachment.key:s}"
|
||||
response = api_client.get("/api/v1.0/chats/media-auth/", HTTP_X_ORIGINAL_URL=original_url)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert "Authorization" not in response
|
||||
|
||||
|
||||
def test_api_media_auth_owner_missing_attachment(api_client):
|
||||
"""
|
||||
Owners of a conversation should not be able to retrieve attachments
|
||||
that are not on the storage.
|
||||
"""
|
||||
attachment = ChatConversationAttachmentFactory()
|
||||
|
||||
original_url = f"http://localhost/media/{attachment.key:s}"
|
||||
now = timezone.now()
|
||||
with freeze_time(now):
|
||||
api_client.force_login(attachment.conversation.owner)
|
||||
response = api_client.get("/api/v1.0/chats/media-auth/", HTTP_X_ORIGINAL_URL=original_url)
|
||||
|
||||
assert response.status_code == 403
|
||||
assert "Authorization" not in response
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
import pytest
|
||||
from rest_framework import status
|
||||
from rest_framework.exceptions import ErrorDetail
|
||||
|
||||
from core.factories import UserFactory
|
||||
|
||||
@@ -27,34 +26,6 @@ def test_update_conversation(api_client):
|
||||
# Verify in database
|
||||
conversation = ChatConversation.objects.get(id=chat_conversation.pk)
|
||||
assert conversation.title == "Updated Title"
|
||||
assert conversation.title_set_by_user_at is not None
|
||||
|
||||
|
||||
def test_update_conversation_limit_title_length(api_client):
|
||||
"""Test that updating a conversation with a title exceeding 100 characters fails validation."""
|
||||
chat_conversation = ChatConversationFactory(title="Initial title")
|
||||
|
||||
url = f"/api/v1.0/chats/{chat_conversation.pk}/"
|
||||
# Create a 101-character title to exceed the 100-character maximum limit
|
||||
new_title = "X" * 101
|
||||
data = {"title": new_title}
|
||||
api_client.force_login(chat_conversation.owner)
|
||||
response = api_client.put(url, data, format="json")
|
||||
|
||||
assert response.status_code == status.HTTP_400_BAD_REQUEST
|
||||
|
||||
assert response.data == {
|
||||
"title": [
|
||||
ErrorDetail(
|
||||
string="Ensure this field has no more than 100 characters.", code="max_length"
|
||||
)
|
||||
]
|
||||
}
|
||||
|
||||
# Verify in database (title should remain unchanged)
|
||||
conversation = ChatConversation.objects.get(id=chat_conversation.pk)
|
||||
assert conversation.title == "Initial title"
|
||||
assert not conversation.title_set_by_user_at
|
||||
|
||||
|
||||
def test_update_conversation_anonymous(api_client):
|
||||
|
||||
@@ -18,28 +18,18 @@ def get_pydantic_tools_by_name(name: str) -> Tool:
|
||||
tool_dict = {
|
||||
"get_current_weather": Tool(get_current_weather, takes_ctx=False),
|
||||
"web_search_brave": Tool(
|
||||
web_search_brave,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
web_search_brave, takes_ctx=False, prepare=only_if_web_search_enabled
|
||||
),
|
||||
"web_search_brave_with_document_backend": Tool(
|
||||
web_search_brave_with_document_backend,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
),
|
||||
"web_search_tavily": Tool(
|
||||
web_search_tavily,
|
||||
takes_ctx=False,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
web_search_tavily, takes_ctx=False, prepare=only_if_web_search_enabled
|
||||
),
|
||||
"web_search_albert_rag": Tool(
|
||||
web_search_albert_rag,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
web_search_albert_rag, takes_ctx=True, prepare=only_if_web_search_enabled
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@@ -1,137 +0,0 @@
|
||||
"""
|
||||
Helpers to add RAG document search tools to an agent based on settings.
|
||||
|
||||
The purpose is to provide a generic way to add multiple RAG document search tools
|
||||
to an agent based on configuration in settings. Each tool can target specific
|
||||
document collections and have its own description.
|
||||
|
||||
Our use case implies that different users might have access to different document collections,
|
||||
so the tools added to the agent are also user-specific.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.contrib.auth import get_user_model
|
||||
from django.utils.module_loading import import_string
|
||||
|
||||
from httpx import HTTPStatusError
|
||||
from pydantic_ai import Agent, ModelRetry, RunContext, RunUsage
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from core.feature_flags.helpers import is_feature_enabled
|
||||
|
||||
from chat.tools.utils import last_model_retry_soft_fail
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
User = get_user_model()
|
||||
|
||||
|
||||
def get_specific_rag_search_tool_config(user: User) -> dict:
|
||||
"""
|
||||
Get the specific RAG search tool configuration from settings.
|
||||
|
||||
Settings example:
|
||||
SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {
|
||||
"french_public_services": {
|
||||
"collection_ids": [784, 785],
|
||||
"feature_flag_value": "disabled",
|
||||
"tool_description": (
|
||||
"Use this tool when the user asks for information about French public services, "
|
||||
"the French labor market, employment laws, social benefits, or "
|
||||
"assistance with administrative procedures."
|
||||
),
|
||||
},
|
||||
}
|
||||
"""
|
||||
return {
|
||||
tool_name: tool_config
|
||||
for tool_name, tool_config in settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS.items()
|
||||
if is_feature_enabled(user, tool_name)
|
||||
}
|
||||
|
||||
|
||||
def _create_document_search_rag(agent, name, description, backend, ids):
|
||||
"""Factory function to create a document search RAG tool."""
|
||||
|
||||
@agent.tool(
|
||||
name=name,
|
||||
retries=1,
|
||||
require_parameter_descriptions=True,
|
||||
description=description,
|
||||
)
|
||||
@last_model_retry_soft_fail
|
||||
async def document_search_rag(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The query to search information about.
|
||||
"""
|
||||
document_store = backend(read_only_collection_id=ids)
|
||||
|
||||
try:
|
||||
rag_results = await document_store.asearch(query)
|
||||
except HTTPStatusError as exc:
|
||||
logger.error(
|
||||
"RAG document search failed for tool %s with error: %s", name, exc, exc_info=True
|
||||
)
|
||||
raise ModelRetry(f"Document search service is currently unavailable: {exc}") from exc
|
||||
|
||||
ctx.usage += RunUsage(
|
||||
input_tokens=rag_results.usage.prompt_tokens,
|
||||
output_tokens=rag_results.usage.completion_tokens,
|
||||
)
|
||||
|
||||
return ToolReturn(
|
||||
return_value={
|
||||
str(idx): {
|
||||
"url": result.url,
|
||||
"snippets": result.content,
|
||||
}
|
||||
for idx, result in enumerate(rag_results.data)
|
||||
},
|
||||
metadata={"sources": {result.url for result in rag_results.data}},
|
||||
)
|
||||
|
||||
return document_search_rag
|
||||
|
||||
|
||||
def add_document_rag_search_tool_from_setting(agent: Agent, user: User) -> None:
|
||||
"""
|
||||
This function takes a configuration setting and generates specific search RAG tools and add
|
||||
it to the agent.
|
||||
|
||||
Args:
|
||||
agent (Agent): The agent to which the tool will be added.
|
||||
user (User): The user for whom the tool is being added.
|
||||
"""
|
||||
|
||||
for tool_name, tool_config in get_specific_rag_search_tool_config(user).items():
|
||||
document_store_backend_name = tool_config.get(
|
||||
"rag_backend_name", settings.RAG_DOCUMENT_SEARCH_BACKEND
|
||||
)
|
||||
try:
|
||||
document_store_backend = import_string(document_store_backend_name)
|
||||
except ImportError as exc:
|
||||
logger.warning(
|
||||
"Could not import RAG backend %s: %s",
|
||||
document_store_backend_name,
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
continue # Skip if the backend is not available
|
||||
|
||||
collection_ids = tool_config.get("collection_ids", [])
|
||||
if not collection_ids:
|
||||
logger.warning("No collection IDs provided for tool %s, skipping.", tool_name)
|
||||
continue # Skip if no collection IDs are provided
|
||||
|
||||
tool_description = tool_config.get("tool_description")
|
||||
if not tool_description:
|
||||
logger.warning("No tool description provided for tool %s, skipping.", tool_name)
|
||||
continue # Skip if no tool description is provided
|
||||
|
||||
_create_document_search_rag(
|
||||
agent, tool_name, tool_description, document_store_backend, collection_ids
|
||||
)
|
||||
@@ -20,7 +20,7 @@ def add_document_rag_search_tool(agent: Agent) -> None:
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The query to search the documents for.
|
||||
query (str): The term to search the internet for.
|
||||
"""
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
|
||||
@@ -43,6 +43,8 @@ def add_document_rag_search_tool(agent: Agent) -> None:
|
||||
def document_rag_instructions() -> str:
|
||||
"""Dynamic system prompt function to add RAG instructions if any."""
|
||||
return (
|
||||
"Use document_search_rag ONLY to retrieve specific passages from attached documents. "
|
||||
"Do NOT use it to summarize; for summaries, call the summarize tool instead."
|
||||
"If the user wants specific information from a document, invoke "
|
||||
"web_search_albert_rag with an appropriate query string."
|
||||
"Do not ask the user for the document; rely on the tool to locate "
|
||||
"and return relevant passages."
|
||||
)
|
||||
|
||||
@@ -1,189 +0,0 @@
|
||||
"""Summarization tool used for uploaded documents."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
import semchunk
|
||||
from asgiref.sync import sync_to_async
|
||||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from chat.agents.summarize import SummarizationAgent
|
||||
from chat.tools.exceptions import ModelCannotRetry
|
||||
from chat.tools.utils import last_model_retry_soft_fail
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def read_document_content(doc):
|
||||
"""Read document content asynchronously."""
|
||||
with default_storage.open(doc.key) as f:
|
||||
return doc.file_name, f.read().decode("utf-8")
|
||||
|
||||
|
||||
async def summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx):
|
||||
"""Summarize a single chunk of text."""
|
||||
sum_prompt = (
|
||||
"You are an agent specializing in text summarization. "
|
||||
"Generate a clear and concise summary of the following passage "
|
||||
f"(part {idx}/{total_chunks}):\n'''\n{chunk}\n'''\n\n"
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"[summarize] CHUNK %s/%s prompt=> %s", idx, total_chunks, sum_prompt[0:100] + "..."
|
||||
)
|
||||
|
||||
try:
|
||||
resp = await summarization_agent.run(sum_prompt, usage=ctx.usage)
|
||||
except Exception as exc:
|
||||
logger.warning("Error during chunk summarization: %s", exc, exc_info=True)
|
||||
raise ModelRetry(
|
||||
"An error occurred while summarizing a part of the document chunk."
|
||||
) from exc
|
||||
|
||||
logger.debug("[summarize] CHUNK %s/%s response<= %s", idx, total_chunks, resp.output or "")
|
||||
return resp.output or ""
|
||||
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def document_summarize( # pylint: disable=too-many-locals
|
||||
ctx: RunContext, *, instructions: str | None = None
|
||||
) -> ToolReturn:
|
||||
"""
|
||||
Generate a complete, ready-to-use summary of the documents in context
|
||||
(do not request the documents to the user).
|
||||
Return this summary directly to the user WITHOUT any modification,
|
||||
or additional summarization.
|
||||
The summary is already optimized and MUST be presented as-is in the final response
|
||||
or translated preserving the information.
|
||||
|
||||
Instructions are optional but should reflect the user's request.
|
||||
|
||||
Examples:
|
||||
"Summarize this doc in 2 paragraphs" -> instructions = "summary in 2 paragraphs"
|
||||
"Summarize this doc in English" -> instructions = "In English"
|
||||
"Summarize this doc" -> instructions = "" (default)
|
||||
|
||||
Args:
|
||||
instructions (str | None): The instructions the user gave to use for the summarization
|
||||
"""
|
||||
try:
|
||||
instructions_hint = (
|
||||
instructions.strip() if instructions else "The summary should contain 2 or 3 parts."
|
||||
)
|
||||
summarization_agent = SummarizationAgent()
|
||||
|
||||
# Collect documents content
|
||||
text_attachment = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
content_type__startswith="text/",
|
||||
)
|
||||
)
|
||||
|
||||
if not text_attachment:
|
||||
raise ModelCannotRetry(
|
||||
"No text documents found in the conversation. "
|
||||
"You must explain this to the user and ask them to provide documents."
|
||||
)
|
||||
|
||||
documents = [await read_document_content(doc) for doc in text_attachment]
|
||||
|
||||
# Chunk documents and summarize each chunk
|
||||
chunk_size = settings.SUMMARIZATION_CHUNK_SIZE
|
||||
chunker = semchunk.chunkerify(
|
||||
tokenizer_or_token_counter=lambda text: len(text.split()),
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
documents_chunks = chunker(
|
||||
[doc[1] for doc in documents],
|
||||
overlap=settings.SUMMARIZATION_OVERLAP_SIZE,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[summarize] chunking: %s parts (size~%s), instructions='%s'",
|
||||
sum(len(chunks) for chunks in documents_chunks),
|
||||
chunk_size,
|
||||
instructions_hint,
|
||||
)
|
||||
|
||||
# Parallelize the chunk summarization with a semaphore to limit concurrent tasks
|
||||
# because it can be very resource intensive on the LLM backend
|
||||
semaphore = asyncio.Semaphore(settings.SUMMARIZATION_CONCURRENT_REQUESTS)
|
||||
|
||||
async def summarize_chunk_with_semaphore(idx, chunk, total_chunks):
|
||||
"""Summarize a chunk with semaphore-controlled concurrency."""
|
||||
async with semaphore:
|
||||
return await summarize_chunk(idx, chunk, total_chunks, summarization_agent, ctx)
|
||||
|
||||
doc_chunk_summaries = []
|
||||
try:
|
||||
for doc_chunks in documents_chunks:
|
||||
summarization_tasks = [
|
||||
summarize_chunk_with_semaphore(idx, chunk, len(doc_chunks))
|
||||
for idx, chunk in enumerate(doc_chunks, start=1)
|
||||
]
|
||||
chunk_summaries = await asyncio.gather(*summarization_tasks)
|
||||
doc_chunk_summaries.append(chunk_summaries)
|
||||
except ModelRetry as exc:
|
||||
logger.warning("Retryable error during chunk summarization: %s", exc, exc_info=True)
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.warning("Error during chunk summarization: %s", exc, exc_info=True)
|
||||
raise ModelRetry("An error occurred while processing document chunks.") from exc
|
||||
|
||||
context = "\n\n".join(
|
||||
doc_name + "\n\n" + "\n\n".join(summaries)
|
||||
for doc_name, summaries in zip(
|
||||
(doc[0] for doc in documents),
|
||||
doc_chunk_summaries,
|
||||
strict=True,
|
||||
)
|
||||
)
|
||||
|
||||
# Merge chunk summaries into a single concise summary
|
||||
merged_prompt = (
|
||||
"Produce a coherent synthesis from the summaries below.\n\n"
|
||||
f"'''\n{context}\n'''\n\n"
|
||||
"Constraints:\n"
|
||||
"- Summarize without repetition.\n"
|
||||
"- Harmonize style and terminology.\n"
|
||||
"- The final summary must be well-structured and formatted in markdown.\n"
|
||||
f"- Follow the instructions: {instructions_hint}\n"
|
||||
"Respond directly with the final summary."
|
||||
)
|
||||
|
||||
logger.debug("[summarize] MERGE prompt=> %s", merged_prompt)
|
||||
|
||||
try:
|
||||
merged_resp = await summarization_agent.run(merged_prompt, usage=ctx.usage)
|
||||
except Exception as exc:
|
||||
logger.warning("Error during merge summarization: %s", exc, exc_info=True)
|
||||
raise ModelRetry("An error occurred while generating the final summary.") from exc
|
||||
|
||||
final_summary = (merged_resp.output or "").strip()
|
||||
|
||||
if not final_summary:
|
||||
raise ModelRetry("The summarization produced an empty result.")
|
||||
|
||||
logger.debug("[summarize] MERGE response<= %s", final_summary)
|
||||
|
||||
return ToolReturn(
|
||||
return_value=final_summary,
|
||||
metadata={"sources": {doc[0] for doc in documents}},
|
||||
)
|
||||
|
||||
except (ModelCannotRetry, ModelRetry):
|
||||
# Re-raise these as-is
|
||||
raise
|
||||
except Exception as exc:
|
||||
# Unexpected error - stop and inform user
|
||||
logger.exception("Unexpected error in document_summarize: %s", exc)
|
||||
raise ModelCannotRetry(
|
||||
f"An unexpected error occurred during document summarization: {type(exc).__name__}. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from exc
|
||||
@@ -1,23 +0,0 @@
|
||||
"""Exceptions for tool function retries."""
|
||||
|
||||
from pydantic_ai import ModelRetry
|
||||
|
||||
|
||||
class ModelRetryLast(ModelRetry):
|
||||
"""
|
||||
Same as ModelRetry but also holds the last retry message to return when all attempts failed.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, last_retry_message: str):
|
||||
"""Initialize ModelRetryLast with message and last retry message."""
|
||||
self.last_retry_message = last_retry_message
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ModelCannotRetry(ModelRetry):
|
||||
"""
|
||||
Exception to raise when a tool function cannot be retried.
|
||||
|
||||
We use this exception to signal that the model should not attempt to retry
|
||||
the tool call, typically because the error is not transient or recoverable.
|
||||
"""
|
||||
@@ -1,50 +0,0 @@
|
||||
"""Tool calling utilities for the chat agent."""
|
||||
|
||||
import functools
|
||||
import logging
|
||||
from typing import Any, Callable
|
||||
|
||||
from pydantic_ai import ModelRetry, RunContext
|
||||
|
||||
from chat.tools.exceptions import ModelCannotRetry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def last_model_retry_soft_fail(
|
||||
tool_func: Callable[..., Any],
|
||||
) -> Callable[..., Any]:
|
||||
"""
|
||||
Wrap a tool function to handle ModelRetry exceptions.
|
||||
|
||||
If the tool function raises ModelRetry and the maximum number of retries
|
||||
has been reached, a ModelCannotRetry exception is raised instead.
|
||||
|
||||
Args:
|
||||
tool_func: The original tool function to wrap.
|
||||
|
||||
Returns:
|
||||
A wrapped tool function with retry handling.
|
||||
"""
|
||||
|
||||
@functools.wraps(tool_func)
|
||||
async def wrapper(ctx: RunContext, *args, **kwargs) -> Any:
|
||||
try:
|
||||
return await tool_func(ctx, *args, **kwargs)
|
||||
except ModelCannotRetry as exc:
|
||||
return str(exc.message)
|
||||
except ModelRetry as exc:
|
||||
logger.error("Tool '%s' raised ModelRetry: %s", ctx, exc.message)
|
||||
if (ctx.retries.get(ctx.tool_name, 0) + 1) >= ctx.max_retries:
|
||||
logger.error("Max retries reached for tool '%s'.", ctx.tool_name)
|
||||
# A bit of a hack to signal that we cannot retry here, while preventing
|
||||
# the LLM to generate an outdated answer.
|
||||
# We may define a more specific exception later base on ModelRetry which
|
||||
# adds a specific message for this case.
|
||||
return (
|
||||
f"{exc.message} You must explain this to the user and "
|
||||
"not try to answer based on your knowledge."
|
||||
)
|
||||
raise # Re-raise to allow retrying
|
||||
|
||||
return wrapper
|
||||
@@ -1,42 +1,24 @@
|
||||
"""Web search tool using Brave for the chat agent."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import List
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.cache import cache
|
||||
from django.utils.module_loading import import_string
|
||||
from django.utils.text import slugify
|
||||
|
||||
import httpx
|
||||
from asgiref.sync import sync_to_async
|
||||
import requests
|
||||
from pydantic_ai import RunContext, RunUsage
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
from trafilatura import extract
|
||||
from trafilatura import extract, fetch_url
|
||||
from trafilatura.meta import reset_caches
|
||||
|
||||
from chat.tools.exceptions import ModelCannotRetry
|
||||
from chat.tools.utils import last_model_retry_soft_fail
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WebSearchError(Exception):
|
||||
"""Base exception for web search errors."""
|
||||
|
||||
|
||||
class BraveAPIError(WebSearchError):
|
||||
"""Error when calling Brave API."""
|
||||
|
||||
|
||||
class DocumentFetchError(WebSearchError):
|
||||
"""Error when fetching or extracting documents."""
|
||||
|
||||
|
||||
async def llm_summarize_async(query: str, text: str) -> str:
|
||||
def llm_summarize(query: str, text: str) -> str:
|
||||
"""
|
||||
Summarize the text using the LLM summarization agent.
|
||||
|
||||
@@ -51,7 +33,7 @@ async def llm_summarize_async(query: str, text: str) -> str:
|
||||
prompt = f"""
|
||||
Based on the following request, summarize the following text in a concise manner,
|
||||
focusing on the key points regarding the user request.
|
||||
The result should be up to 30 lines long.
|
||||
he result should be up to 30 lines long.
|
||||
|
||||
<user request>
|
||||
{query}
|
||||
@@ -62,87 +44,54 @@ The result should be up to 30 lines long.
|
||||
</text to summarize>
|
||||
"""
|
||||
|
||||
result = await summarization_agent.run(prompt)
|
||||
result = summarization_agent.run_sync(prompt)
|
||||
return result.output
|
||||
|
||||
|
||||
async def _fetch_url_async(url: str, timeout: int = 30) -> str:
|
||||
"""Fetch URL content asynchronously."""
|
||||
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True) as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
return response.text
|
||||
def _fetch_and_extract(url: str) -> str:
|
||||
"""Fetch and extract text content from the URL."""
|
||||
cache_key = f"web_search_brave:extract:{url}"
|
||||
|
||||
|
||||
async def _fetch_and_extract_async(url: str) -> str:
|
||||
"""Fetch and extract text content from the URL asynchronously."""
|
||||
cache_key = f"web_search_brave:extract:{slugify(url)}"
|
||||
|
||||
# Check cache first
|
||||
if (document := await cache.aget(cache_key)) is not None:
|
||||
if (document := cache.get(cache_key)) is not None:
|
||||
return document
|
||||
|
||||
try:
|
||||
# Fetch HTML
|
||||
html = await _fetch_url_async(url, timeout=settings.BRAVE_API_TIMEOUT)
|
||||
html = fetch_url(url)
|
||||
document = extract(html, include_comments=False, no_fallback=True) or ""
|
||||
cache.set(cache_key, document, settings.BRAVE_CACHE_TTL)
|
||||
|
||||
# Extract text in thread pool (trafilatura is CPU-bound)
|
||||
document = await sync_to_async(extract)(html, include_comments=False, no_fallback=True)
|
||||
|
||||
# Cache the result
|
||||
await cache.aset(cache_key, document, settings.BRAVE_CACHE_TTL)
|
||||
return document
|
||||
|
||||
except httpx.HTTPError as e:
|
||||
logger.warning("HTTP error fetching %s: %s", url, e, exc_info=True)
|
||||
raise DocumentFetchError(f"Failed to fetch {url}: {e}") from e
|
||||
except Exception as e:
|
||||
logger.warning("Error extracting content from %s: %s", url, e, exc_info=True)
|
||||
raise DocumentFetchError(f"Failed to extract content from {url}: {e}") from e
|
||||
return document
|
||||
|
||||
|
||||
async def _extract_and_summarize_snippets_async(query: str, url: str) -> List[str]:
|
||||
def _extract_and_summarize_snippets(query: str, url: str) -> List[str]:
|
||||
"""Fetch, extract and summarize text content from the URL.
|
||||
|
||||
Returns a list of snippets (0 or 1 element, preserving existing behavior).
|
||||
"""
|
||||
try:
|
||||
document = await _fetch_and_extract_async(url)
|
||||
if not document:
|
||||
return []
|
||||
|
||||
if not settings.BRAVE_SUMMARIZATION_ENABLED:
|
||||
return [document]
|
||||
|
||||
try:
|
||||
snippet = await llm_summarize_async(query, document)
|
||||
return [snippet] if snippet else []
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
logger.exception("Summarization failed for %s: %s", url, e)
|
||||
# Fallback to raw document if summarization fails
|
||||
return [document]
|
||||
|
||||
except DocumentFetchError:
|
||||
# Document fetch failed, return empty
|
||||
# Cache by URL to avoid repeated fetch/extract across calls
|
||||
document = _fetch_and_extract(url)
|
||||
if not document:
|
||||
return []
|
||||
|
||||
|
||||
async def _fetch_and_store_async(url: str, document_store) -> None:
|
||||
"""Fetch, extract and store text content from the URL in the document store."""
|
||||
if not settings.BRAVE_SUMMARIZATION_ENABLED:
|
||||
return [document]
|
||||
|
||||
try:
|
||||
document = await _fetch_and_extract_async(url)
|
||||
snippet = llm_summarize(query, document)
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
logger.exception("Summarization failed for %s: %s", url, e)
|
||||
snippet = None
|
||||
|
||||
logger.debug("Fetched document: %s", document)
|
||||
|
||||
if document:
|
||||
await document_store.astore_document(url, document)
|
||||
except DocumentFetchError as e:
|
||||
logger.warning("Failed to fetch and store %s: %s", url, e)
|
||||
# Continue with other documents
|
||||
return [snippet] if snippet else []
|
||||
|
||||
|
||||
async def _query_brave_api_async(query: str) -> List[dict]:
|
||||
def _fetch_and_store(url: str, document_store) -> None:
|
||||
"""Fetch, extract and store text content from the URL in the document store."""
|
||||
document = _fetch_and_extract(url)
|
||||
if document:
|
||||
document_store.store_document(url, document)
|
||||
|
||||
|
||||
def _query_brave_api(query: str) -> List[dict]:
|
||||
"""Query the Brave Search API and return the raw results."""
|
||||
url = "https://api.search.brave.com/res/v1/web/search"
|
||||
headers = {
|
||||
@@ -160,53 +109,14 @@ async def _query_brave_api_async(query: str) -> List[dict]:
|
||||
"extra_snippets": settings.BRAVE_SEARCH_EXTRA_SNIPPETS,
|
||||
}
|
||||
params = {k: v for k, v in data.items() if v is not None}
|
||||
response = requests.get(url, headers=headers, params=params, timeout=settings.BRAVE_API_TIMEOUT)
|
||||
response.raise_for_status()
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.BRAVE_API_TIMEOUT) as client:
|
||||
response = await client.get(url, headers=headers, params=params)
|
||||
response.raise_for_status()
|
||||
json_response = response.json()
|
||||
json_response = response.json()
|
||||
|
||||
# https://api-dashboard.search.brave.com/app/documentation/web-search/responses#Result
|
||||
return json_response.get("web", {}).get("results", [])
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 429:
|
||||
# Rate limit - retryable
|
||||
logger.warning("Brave API rate limited: %s", e)
|
||||
raise ModelRetry(
|
||||
"The search API is rate limited. Please wait a moment and try again."
|
||||
) from e
|
||||
if e.response.status_code >= 500:
|
||||
# Server error - retryable
|
||||
logger.warning("Brave API error: %s", e)
|
||||
raise ModelRetry(
|
||||
"The search service is temporarily unavailable due to a server error. Retrying..."
|
||||
) from e
|
||||
|
||||
# Client error (4xx) - not retryable, stop and inform user
|
||||
logger.error("Brave API client error: %s", e)
|
||||
raise ModelCannotRetry(
|
||||
f"Web search failed with a client error (status {e.response.status_code}). "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from e
|
||||
except httpx.TimeoutException as e:
|
||||
# Timeout - retryable
|
||||
logger.warning("Brave API timeout: %s", e)
|
||||
raise ModelRetry("The search request timed out. Retrying with a fresh attempt...") from e
|
||||
except httpx.HTTPError as e:
|
||||
# Other HTTP errors - retryable
|
||||
logger.warning("Brave API connection error: %s", e)
|
||||
raise ModelRetry(
|
||||
f"Connection error while searching the web: {type(e).__name__}. Retrying..."
|
||||
) from e
|
||||
except Exception as e:
|
||||
# Unexpected errors - not retryable, stop completely
|
||||
logger.exception("Unexpected error querying Brave API: %s", e)
|
||||
raise ModelCannotRetry(
|
||||
f"An unexpected error occurred with the search service: {type(e).__name__}. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from e
|
||||
# See https://api-dashboard.search.brave.com/app/documentation/web-search/responses#Result
|
||||
# & https://api-dashboard.search.brave.com/app/documentation/web-search/responses#SearchResult
|
||||
return json_response.get("web", {}).get("results", [])
|
||||
|
||||
|
||||
def format_tool_return(raw_search_results: List[dict]) -> ToolReturn:
|
||||
@@ -230,132 +140,89 @@ def format_tool_return(raw_search_results: List[dict]) -> ToolReturn:
|
||||
)
|
||||
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def web_search_brave(_ctx: RunContext, query: str) -> ToolReturn:
|
||||
def web_search_brave(query: str) -> ToolReturn:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
_ctx (RunContext): The run context, used by the wrapper.
|
||||
query (str): The query to search for.
|
||||
"""
|
||||
try:
|
||||
raw_search_results = await _query_brave_api_async(query)
|
||||
raw_search_results = _query_brave_api(query)
|
||||
|
||||
await sync_to_async(reset_caches)() # Clear trafilatura caches to avoid memory bloat/leaks
|
||||
reset_caches() # Clear trafilatura caches to avoid memory bloat/leaks
|
||||
|
||||
# Parallelize fetch/extract for results that don't include extra_snippets
|
||||
to_process = [
|
||||
(idx, r) for idx, r in enumerate(raw_search_results) if not r.get("extra_snippets")
|
||||
]
|
||||
# Parallelize fetch/extract for results that don't include extra_snippets
|
||||
to_process = [
|
||||
(idx, r) for idx, r in enumerate(raw_search_results) if not r.get("extra_snippets")
|
||||
]
|
||||
|
||||
if to_process:
|
||||
# Process all URLs concurrently
|
||||
tasks = [
|
||||
_extract_and_summarize_snippets_async(query, r["url"]) for idx, r in to_process
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=False)
|
||||
if to_process:
|
||||
max_workers = min(settings.BRAVE_MAX_WORKERS, len(to_process))
|
||||
if max_workers == 1:
|
||||
# Avoid overhead of ThreadPoolExecutor if only one task
|
||||
for idx, r in to_process:
|
||||
raw_search_results[idx]["extra_snippets"] = _extract_and_summarize_snippets(
|
||||
query, r["url"]
|
||||
)
|
||||
|
||||
# Update raw_search_results with extracted snippets
|
||||
for (idx, _), snippets in zip(to_process, results, strict=True):
|
||||
raw_search_results[idx]["extra_snippets"] = snippets
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_map = {
|
||||
executor.submit(_extract_and_summarize_snippets, query, r["url"]): idx
|
||||
for idx, r in to_process
|
||||
}
|
||||
for future in as_completed(future_map):
|
||||
idx = future_map[future]
|
||||
raw_search_results[idx]["extra_snippets"] = future.result()
|
||||
|
||||
formatted_result = format_tool_return(raw_search_results)
|
||||
|
||||
# Check if we got any valid results
|
||||
if not formatted_result.return_value:
|
||||
raise ModelRetry(
|
||||
"No valid search results were extracted from the web pages. "
|
||||
"Retrying the search to find better sources..."
|
||||
)
|
||||
|
||||
return formatted_result
|
||||
|
||||
except (ModelCannotRetry, ModelRetry):
|
||||
# Re-raise these as-is
|
||||
raise
|
||||
except Exception as exc:
|
||||
# Unexpected error in our code - stop and inform user
|
||||
logger.exception("Unexpected error in web_search_brave: %s", exc)
|
||||
raise ModelCannotRetry(
|
||||
f"An unexpected error occurred during web search: {type(exc).__name__}. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from exc
|
||||
return format_tool_return(raw_search_results)
|
||||
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn:
|
||||
def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Search the web for up-to-date information using RAG backend
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The query to search for.
|
||||
"""
|
||||
logger.info("Starting web search with RAG backend for query: %s", query)
|
||||
try:
|
||||
raw_search_results = await _query_brave_api_async(query)
|
||||
raw_search_results = _query_brave_api(query)
|
||||
|
||||
# Clear trafilatura caches in thread pool to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, reset_caches)
|
||||
reset_caches() # Clear trafilatura caches to avoid memory bloat/leaks
|
||||
|
||||
# Store documents in a temporary document store for RAG search
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
|
||||
# Create temporary collection
|
||||
temp_collection_name = f"tmp-{uuid.uuid4()}"
|
||||
try:
|
||||
async with document_store_backend.temporary_collection_async(
|
||||
temp_collection_name
|
||||
) as document_store:
|
||||
# Fetch and store all documents concurrently
|
||||
tasks = [
|
||||
_fetch_and_store_async(result["url"], document_store)
|
||||
# Store documents in a temporary document store for RAG search
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
with document_store_backend.temporary_collection(f"tmp-{uuid.uuid4()}") as document_store:
|
||||
max_workers = min(settings.BRAVE_MAX_WORKERS, len(raw_search_results))
|
||||
if max_workers == 1:
|
||||
for result in raw_search_results:
|
||||
# Fetch and extract document content
|
||||
_fetch_and_store(result["url"], document_store)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(_fetch_and_store, result["url"], document_store)
|
||||
for result in raw_search_results
|
||||
]
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
future.result()
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
logger.exception("Error fetching/storing document: %s", e)
|
||||
|
||||
# Perform RAG search
|
||||
rag_results = await document_store.asearch(
|
||||
query,
|
||||
results_count=settings.BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER,
|
||||
)
|
||||
logger.info("RAG search returned: %s", rag_results)
|
||||
rag_results = document_store.search(query)
|
||||
|
||||
ctx.usage += RunUsage(
|
||||
input_tokens=rag_results.usage.prompt_tokens,
|
||||
output_tokens=rag_results.usage.completion_tokens,
|
||||
)
|
||||
ctx.usage += RunUsage(
|
||||
input_tokens=rag_results.usage.prompt_tokens,
|
||||
output_tokens=rag_results.usage.completion_tokens,
|
||||
)
|
||||
|
||||
# Map RAG results back to raw search results to include extra_snippets
|
||||
for rag_result in rag_results.data:
|
||||
for result in raw_search_results:
|
||||
if result["url"] == rag_result.url:
|
||||
result.setdefault("extra_snippets", []).append(rag_result.content)
|
||||
break
|
||||
# Map RAG results back to raw search results to include extra_snippets
|
||||
# Suboptimal O(N^2) but N is small...
|
||||
for rag_result in rag_results.data:
|
||||
for result in raw_search_results:
|
||||
if result["url"] == rag_result.url:
|
||||
result.setdefault("extra_snippets", []).append(rag_result.content)
|
||||
break
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("Error with document store: %s", exc)
|
||||
raise ModelRetry(
|
||||
f"Document storage temporarily failed: {type(exc).__name__}. "
|
||||
"Retrying the operation..."
|
||||
) from exc
|
||||
|
||||
formatted_result = format_tool_return(raw_search_results)
|
||||
|
||||
# Check if we got any valid results
|
||||
if not formatted_result.return_value:
|
||||
raise ModelRetry("No valid search results were extracted.")
|
||||
|
||||
return formatted_result
|
||||
except (ModelCannotRetry, ModelRetry):
|
||||
# Re-raise these as-is
|
||||
raise
|
||||
except Exception as e:
|
||||
# Unexpected error - stop and inform user
|
||||
logger.exception("Unexpected error in web_search_brave_with_document_backend: %s", e)
|
||||
raise ModelCannotRetry(
|
||||
f"An unexpected error occurred during web search with RAG: {type(e).__name__}. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from e
|
||||
return format_tool_return(raw_search_results)
|
||||
|
||||
+8
-197
@@ -1,26 +1,17 @@
|
||||
"""Chat API implementation."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from uuid import uuid4
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
from django.http import Http404, StreamingHttpResponse
|
||||
from django.http import StreamingHttpResponse
|
||||
|
||||
import langfuse
|
||||
import magic
|
||||
import posthog
|
||||
from lasuite.malware_detection import malware_detection
|
||||
from rest_framework import decorators, filters, mixins, permissions, status, viewsets
|
||||
from rest_framework.exceptions import MethodNotAllowed, PermissionDenied, ValidationError
|
||||
from rest_framework.exceptions import ValidationError
|
||||
from rest_framework.response import Response
|
||||
from rest_framework.views import APIView
|
||||
|
||||
from core.api.viewsets import Pagination, SerializerPerActionMixin
|
||||
from core.file_upload import enums
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.mixins import AttachmentMixin
|
||||
from core.filters import remove_accents
|
||||
|
||||
from activation_codes.permissions import IsActivatedUser
|
||||
@@ -41,48 +32,13 @@ class ChatConversationFilter(filters.BaseFilterBackend):
|
||||
return queryset
|
||||
|
||||
|
||||
class ChatAttachmentMixin(AttachmentMixin): # pylint: disable=abstract-method
|
||||
"""Mixin to handle attachment authorization for chat conversations."""
|
||||
|
||||
@decorators.action(detail=True, methods=["post"], url_path="attachment-upload")
|
||||
def attachment_upload(self, request, *args, **kwargs):
|
||||
"""Explicitly disable this action."""
|
||||
raise MethodNotAllowed("POST")
|
||||
|
||||
@decorators.action(detail=True, methods=["get"], url_path="media-check")
|
||||
def media_check(self, request, *args, **kwargs):
|
||||
"""Explicitly disable this action."""
|
||||
raise MethodNotAllowed("GET")
|
||||
|
||||
def check_attachment_holder_permission(self, user, url_params, key):
|
||||
"""
|
||||
Check if the user has permission to access the holder of the attachment.
|
||||
|
||||
Raises PermissionDenied if the user does not have permission.
|
||||
"""
|
||||
if not user.is_authenticated:
|
||||
raise PermissionDenied()
|
||||
|
||||
try:
|
||||
models.ChatConversation.objects.get(
|
||||
pk=url_params["pk"],
|
||||
owner=user,
|
||||
)
|
||||
# We don't need to check the ChatConversationAttachment here because
|
||||
# if the storage object exists, it means the attachment is linked
|
||||
# to the conversation, which is already verified by the above query.
|
||||
except models.ChatConversation.DoesNotExist as exc:
|
||||
raise PermissionDenied() from exc
|
||||
|
||||
|
||||
class ChatViewSet( # pylint: disable=too-many-ancestors, abstract-method
|
||||
class ChatViewSet( # pylint: disable=too-many-ancestors
|
||||
SerializerPerActionMixin,
|
||||
mixins.CreateModelMixin,
|
||||
mixins.RetrieveModelMixin,
|
||||
mixins.ListModelMixin,
|
||||
mixins.DestroyModelMixin,
|
||||
mixins.UpdateModelMixin,
|
||||
ChatAttachmentMixin,
|
||||
viewsets.GenericViewSet,
|
||||
):
|
||||
"""ViewSet for managing chat conversations.
|
||||
@@ -105,22 +61,10 @@ class ChatViewSet( # pylint: disable=too-many-ancestors, abstract-method
|
||||
filter_backends = [filters.OrderingFilter, ChatConversationFilter]
|
||||
ordering = ["-created_at"]
|
||||
ordering_fields = ["created_at", "updated_at"]
|
||||
queryset = models.ChatConversation.objects # defined to be used in AttachmentMixin
|
||||
|
||||
def get_queryset(self):
|
||||
"""Return the queryset for the chat conversations."""
|
||||
return (
|
||||
self.queryset.filter(owner=self.request.user)
|
||||
if self.request.user.is_authenticated
|
||||
else self.queryset.none()
|
||||
)
|
||||
|
||||
def get_permissions(self):
|
||||
"""Return the permissions for the viewset."""
|
||||
if self.action in ["media_auth", "media_check"]:
|
||||
# Permission is checked in AttachmentMixin
|
||||
self.permission_classes = []
|
||||
return super().get_permissions()
|
||||
return models.ChatConversation.objects.filter(owner=self.request.user)
|
||||
|
||||
@decorators.action(
|
||||
methods=["post"],
|
||||
@@ -179,32 +123,10 @@ class ChatViewSet( # pylint: disable=too-many-ancestors, abstract-method
|
||||
or self.request.LANGUAGE_CODE # from the LocaleMiddleware
|
||||
),
|
||||
)
|
||||
|
||||
# This environment variable allows switching between sync and async streaming modes
|
||||
# based on the server configuration. Tests run in sync mode (WSGI), while
|
||||
# production uses async mode (Uvicorn ASGI).
|
||||
is_async_mode = os.environ.get("PYTHON_SERVER_MODE", "sync") == "async"
|
||||
|
||||
if is_async_mode:
|
||||
logger.debug("Using ASYNC streaming for chat conversation.")
|
||||
if protocol == "data":
|
||||
streaming_content = ai_service.stream_data_async(
|
||||
messages, force_web_search=force_web_search
|
||||
)
|
||||
else: # Default to 'text' protocol
|
||||
streaming_content = ai_service.stream_text_async(
|
||||
messages, force_web_search=force_web_search
|
||||
)
|
||||
else:
|
||||
logger.debug("Using SYNC streaming for chat conversation.")
|
||||
if protocol == "data":
|
||||
streaming_content = ai_service.stream_data(
|
||||
messages, force_web_search=force_web_search
|
||||
)
|
||||
else: # Default to 'text' protocol
|
||||
streaming_content = ai_service.stream_text(
|
||||
messages, force_web_search=force_web_search
|
||||
)
|
||||
if protocol == "data":
|
||||
streaming_content = ai_service.stream_data(messages, force_web_search=force_web_search)
|
||||
else: # Default to 'text' protocol
|
||||
streaming_content = ai_service.stream_text(messages, force_web_search=force_web_search)
|
||||
|
||||
response = StreamingHttpResponse(
|
||||
streaming_content,
|
||||
@@ -324,114 +246,3 @@ class LLMConfigurationView(APIView):
|
||||
},
|
||||
)
|
||||
return Response(serializer.data, status=status.HTTP_200_OK)
|
||||
|
||||
|
||||
class ChatConversationAttachmentViewSet(
|
||||
SerializerPerActionMixin,
|
||||
mixins.CreateModelMixin,
|
||||
mixins.RetrieveModelMixin,
|
||||
viewsets.GenericViewSet,
|
||||
):
|
||||
"""ViewSet for managing chat conversation attachments.
|
||||
|
||||
Provides endpoints to create and retrieve chat conversation attachments.
|
||||
"""
|
||||
|
||||
pagination_class = None # No pagination for attachments
|
||||
permission_classes = [
|
||||
IsActivatedUser, # see activation_codes application
|
||||
permissions.IsAuthenticated,
|
||||
]
|
||||
serializer_class = serializers.ChatConversationAttachmentSerializer
|
||||
create_serializer_class = serializers.CreateChatConversationAttachmentSerializer
|
||||
queryset = models.ChatConversationAttachment.objects
|
||||
|
||||
def get_queryset(self):
|
||||
"""Return the queryset for the chat conversation attachments."""
|
||||
return (
|
||||
self.queryset.filter(
|
||||
conversation_id=self.kwargs["conversation_pk"],
|
||||
conversation__owner=self.request.user,
|
||||
)
|
||||
if self.request.user.is_authenticated
|
||||
else self.queryset.none()
|
||||
)
|
||||
|
||||
def get_serializer_context(self):
|
||||
"""Return the context for the serializer."""
|
||||
context = super().get_serializer_context()
|
||||
context["conversation_pk"] = self.kwargs["conversation_pk"]
|
||||
return context
|
||||
|
||||
def perform_create(self, serializer):
|
||||
"""Set the uploaded_by field to the current user."""
|
||||
# assert the user is the owner of the conversation
|
||||
if not models.ChatConversation.objects.filter(
|
||||
pk=self.kwargs["conversation_pk"],
|
||||
owner=self.request.user,
|
||||
).exists():
|
||||
raise Http404
|
||||
|
||||
file_name = serializer.validated_data["file_name"]
|
||||
extension = file_name.rpartition(".")[-1] if "." in file_name else None
|
||||
|
||||
file_id = uuid4()
|
||||
holder_key_base = f"{self.kwargs['conversation_pk']!s}"
|
||||
ext_suffix = f".{extension}" if extension else ""
|
||||
key = f"{holder_key_base}/{AttachmentMixin.ATTACHMENTS_FOLDER:s}/{file_id!s}{ext_suffix}"
|
||||
|
||||
serializer.save(
|
||||
conversation_id=self.kwargs["conversation_pk"],
|
||||
uploaded_by=self.request.user,
|
||||
upload_state=enums.AttachmentStatus.PENDING,
|
||||
key=key,
|
||||
)
|
||||
|
||||
@decorators.action(detail=True, methods=["post"], url_path="upload-ended")
|
||||
def upload_ended(self, request, *args, **kwargs):
|
||||
"""
|
||||
Start the analysis of an item after a successful upload.
|
||||
"""
|
||||
|
||||
attachment = self.get_object()
|
||||
|
||||
if attachment.upload_state != AttachmentStatus.PENDING:
|
||||
raise ValidationError(
|
||||
{"attachment": "This action is only available for items in PENDING state."},
|
||||
code="upload-state-not-pending",
|
||||
)
|
||||
|
||||
mime_detector = magic.Magic(mime=True)
|
||||
with default_storage.open(attachment.key, "rb") as file:
|
||||
mimetype = mime_detector.from_buffer(file.read(2048))
|
||||
size = file.size
|
||||
|
||||
attachment.upload_state = AttachmentStatus.ANALYZING
|
||||
attachment.content_type = mimetype
|
||||
attachment.size = size
|
||||
|
||||
attachment.save(update_fields=["upload_state", "content_type", "size"])
|
||||
|
||||
malware_detection.analyse_file(
|
||||
attachment.key,
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=self.kwargs["conversation_pk"],
|
||||
)
|
||||
|
||||
serializer = self.get_serializer(attachment)
|
||||
|
||||
if settings.POSTHOG_KEY:
|
||||
posthog.capture(
|
||||
"item_uploaded",
|
||||
distinct_id=str(request.user.pk), # same as set by the frontend
|
||||
properties={
|
||||
"id": attachment.pk,
|
||||
"file_name": attachment.file_name,
|
||||
"size": attachment.size,
|
||||
"mimetype": attachment.content_type,
|
||||
},
|
||||
)
|
||||
|
||||
return Response(serializer.data, status=status.HTTP_200_OK)
|
||||
|
||||
+1
-16
@@ -1,6 +1,5 @@
|
||||
"""Global fixtures for the backend tests."""
|
||||
|
||||
import posthog
|
||||
import pytest
|
||||
from rest_framework.test import APIClient
|
||||
from urllib3.connectionpool import HTTPConnectionPool
|
||||
@@ -22,7 +21,7 @@ def no_http_requests(monkeypatch):
|
||||
Credits: https://blog.jerrycodes.com/no-http-requests/
|
||||
"""
|
||||
|
||||
allowed_hosts = {"localhost", "minio", "minio:9000"}
|
||||
allowed_hosts = {"localhost"}
|
||||
original_urlopen = HTTPConnectionPool.urlopen
|
||||
|
||||
def urlopen_mock(self, method, url, *args, **kwargs):
|
||||
@@ -42,17 +41,3 @@ def feature_flags_fixture(settings):
|
||||
"""
|
||||
settings.FEATURE_FLAGS = settings.FEATURE_FLAGS.model_copy(deep=True)
|
||||
yield settings.FEATURE_FLAGS
|
||||
|
||||
|
||||
@pytest.fixture(name="posthog", scope="function")
|
||||
def posthog_fixture(settings):
|
||||
"""Mock PostHog in tests to avoid real network calls."""
|
||||
settings.POSTHOG_KEY = {"id": "132456", "host": "https://eu.i.posthog-test.com"}
|
||||
|
||||
posthog.api_key = settings.POSTHOG_KEY["id"]
|
||||
posthog.host = settings.POSTHOG_KEY["host"]
|
||||
|
||||
yield posthog
|
||||
|
||||
posthog.api_key = None
|
||||
posthog.host = None
|
||||
|
||||
@@ -1,18 +0,0 @@
|
||||
"""
|
||||
ASGI config for conversations project.
|
||||
|
||||
It exposes the ASGI callable as a module-level variable named ``application``.
|
||||
|
||||
For more information on this file, see
|
||||
https://docs.djangoproject.com/en/dev/howto/deployment/asgi/
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
from configurations.asgi import get_asgi_application
|
||||
|
||||
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "conversations.settings")
|
||||
os.environ.setdefault("DJANGO_CONFIGURATION", "Development")
|
||||
os.environ.setdefault("PYTHON_SERVER_MODE", "async")
|
||||
|
||||
application = get_asgi_application()
|
||||
@@ -23,13 +23,6 @@ class BraveSettings:
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
# For web_search_brave_with_document_backend: number of chunks to retrieve RAG search
|
||||
BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER = values.IntegerValue(
|
||||
default=10,
|
||||
environ_name="BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
# For optimal performance, BRAVE_MAX_WORKERS should be equal to BRAVE_MAX_RESULTS
|
||||
# also considering the number of concurrent requests your server can handle.
|
||||
BRAVE_MAX_WORKERS = values.IntegerValue(
|
||||
|
||||
@@ -1,43 +0,0 @@
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"hrid": "default-model",
|
||||
"model_name": "mistral-mock",
|
||||
"human_readable_name": "Default Model",
|
||||
"provider_name": "default-provider",
|
||||
"profile": null,
|
||||
"settings": {},
|
||||
"is_active": true,
|
||||
"icon": [
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABwAAAAcCAMAAABF0y+mAAAAn1BMVEUALosAKoovTZjw8vb////+9/jlPUniAAz",
|
||||
"iABUAGIWbpsTwq7HhAAAAI4dle7DrdX4AJohRaaboXWj7+/zn6On5//9NZaT29vfoWmVHYKDoUl/k5OUAIYddc6vpbHYCM47Y3+v53+LiFCUA",
|
||||
"HIWnsckYPJHi6PL77O7jJjW3wdf1w8jre4QgQ5TZ2txwg7Pr3+I8WZ6OnsTuoamClL7tlZ5xz5y8AAAAzUlEQVR4AZ3RRQKDQBBEUSTu7h5c4",
|
||||
"vc/W6Yp3KG2Dz4ynDdeEBvOmq12xx2E1u0B+4NOEocj4DgNJ1PgLAvni8WyBq5Yc71ubFJx23C2q4P7dRYejg1xzvCUgvz5guz11k7gXYKF/1",
|
||||
"8oyiYuvHAYeVkhXCzolVStHcGDjiQzNmMQxsMI5rEJRdQSPZvbpE2E8aY6gC6Z+2Hg4dFA0Yb4YedNL/v4Fk8WJuwiGhrChJNXI210rnib9Fs",
|
||||
"JlXRUC/HwTscPIXf/iklq/tjb/gHAdxkCUjAg2QAAAABJRU5ErkJggg=="
|
||||
],
|
||||
"system_prompt": "You are a helpful AI assistant.",
|
||||
"tools": []
|
||||
},
|
||||
{
|
||||
"hrid": "default-summarization-model",
|
||||
"model_name": "mistral-mock",
|
||||
"human_readable_name": "Default Summarization Model",
|
||||
"provider_name": "default-provider",
|
||||
"profile": null,
|
||||
"settings": {},
|
||||
"is_active": true,
|
||||
"icon": null,
|
||||
"system_prompt": "You are a helpful AI assistant specialized in summarization.",
|
||||
"tools": []
|
||||
}
|
||||
],
|
||||
"providers": [
|
||||
{
|
||||
"hrid": "default-provider",
|
||||
"base_url": "http://host.docker.internal:8900",
|
||||
"api_key": "openmockllm-api-key",
|
||||
"kind": "mistral"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -109,7 +109,7 @@ class Base(BraveSettings, Configuration):
|
||||
|
||||
STORAGES = {
|
||||
"default": {
|
||||
"BACKEND": "storages.backends.s3.S3Storage",
|
||||
"BACKEND": "django.core.files.storage.FileSystemStorage",
|
||||
},
|
||||
"staticfiles": {
|
||||
"BACKEND": values.Value(
|
||||
@@ -119,132 +119,6 @@ class Base(BraveSettings, Configuration):
|
||||
},
|
||||
}
|
||||
|
||||
# Media
|
||||
AWS_S3_ENDPOINT_URL = values.Value(environ_name="AWS_S3_ENDPOINT_URL", environ_prefix=None)
|
||||
AWS_S3_ACCESS_KEY_ID = values.Value(environ_name="AWS_S3_ACCESS_KEY_ID", environ_prefix=None)
|
||||
AWS_S3_SECRET_ACCESS_KEY = values.Value(
|
||||
environ_name="AWS_S3_SECRET_ACCESS_KEY", environ_prefix=None
|
||||
)
|
||||
AWS_S3_REGION_NAME = values.Value(environ_name="AWS_S3_REGION_NAME", environ_prefix=None)
|
||||
AWS_STORAGE_BUCKET_NAME = values.Value(
|
||||
"conversations-media-storage",
|
||||
environ_name="AWS_STORAGE_BUCKET_NAME",
|
||||
environ_prefix=None,
|
||||
)
|
||||
AWS_S3_SIGNATURE_VERSION = values.Value(
|
||||
"s3v4",
|
||||
environ_name="AWS_S3_SIGNATURE_VERSION",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
AWS_S3_UPLOAD_POLICY_EXPIRATION = values.Value(
|
||||
60, # 1 minute
|
||||
environ_name="AWS_S3_UPLOAD_POLICY_EXPIRATION",
|
||||
environ_prefix=None,
|
||||
)
|
||||
AWS_S3_RETRIEVE_POLICY_EXPIRATION = values.Value(
|
||||
3 * 60, # 3 minutes
|
||||
environ_name="AWS_S3_RETRIEVE_POLICY_EXPIRATION",
|
||||
environ_prefix=None,
|
||||
)
|
||||
AWS_S3_DOMAIN_REPLACE = values.Value(
|
||||
environ_name="AWS_S3_DOMAIN_REPLACE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
ATTACHMENT_CHECK_UNSAFE_MIME_TYPES_ENABLED = values.BooleanValue(
|
||||
True,
|
||||
environ_name="ATTACHMENT_CHECK_UNSAFE_MIME_TYPES_ENABLED",
|
||||
environ_prefix=None,
|
||||
)
|
||||
ATTACHMENT_UNSAFE_MIME_TYPES = [
|
||||
# Executable Files
|
||||
"application/x-msdownload",
|
||||
"application/x-bat",
|
||||
"application/x-dosexec",
|
||||
"application/x-sh",
|
||||
"application/x-ms-dos-executable",
|
||||
"application/x-msi",
|
||||
"application/java-archive",
|
||||
"application/octet-stream",
|
||||
# Dynamic Web Pages
|
||||
"application/x-httpd-php",
|
||||
"application/x-asp",
|
||||
"application/x-aspx",
|
||||
"application/jsp",
|
||||
"application/xhtml+xml",
|
||||
"application/x-python-code",
|
||||
"application/x-perl",
|
||||
"text/html",
|
||||
"text/javascript",
|
||||
"text/x-php",
|
||||
# System Files
|
||||
"application/x-msdownload",
|
||||
"application/x-sys",
|
||||
"application/x-drv",
|
||||
"application/cpl",
|
||||
"application/x-apple-diskimage",
|
||||
# Script Files
|
||||
"application/javascript",
|
||||
"application/x-vbscript",
|
||||
"application/x-powershell",
|
||||
"application/x-shellscript",
|
||||
# Compressed/Archive Files
|
||||
"application/zip",
|
||||
"application/x-tar",
|
||||
"application/gzip",
|
||||
"application/x-bzip2",
|
||||
"application/x-7z-compressed",
|
||||
"application/x-rar",
|
||||
"application/x-rar-compressed",
|
||||
"application/x-compress",
|
||||
"application/x-lzma",
|
||||
# Macros in Documents
|
||||
"application/vnd.ms-word",
|
||||
"application/vnd.ms-excel",
|
||||
"application/vnd.ms-powerpoint",
|
||||
"application/vnd.ms-word.document.macroenabled.12",
|
||||
"application/vnd.ms-excel.sheet.macroenabled.12",
|
||||
"application/vnd.ms-powerpoint.presentation.macroenabled.12",
|
||||
# Disk Images & Virtual Disk Files
|
||||
"application/x-iso9660-image",
|
||||
"application/x-vmdk",
|
||||
"application/x-apple-diskimage",
|
||||
"application/x-dmg",
|
||||
# Other Dangerous MIME Types
|
||||
"application/x-ms-application",
|
||||
"application/x-msdownload",
|
||||
"application/x-shockwave-flash",
|
||||
"application/x-silverlight-app",
|
||||
"application/x-java-vm",
|
||||
"application/x-bittorrent",
|
||||
"application/hta",
|
||||
"application/x-csh",
|
||||
"application/x-ksh",
|
||||
"application/x-ms-regedit",
|
||||
"application/x-msdownload",
|
||||
"application/xml",
|
||||
]
|
||||
ATTACHMENT_MAX_SIZE = values.IntegerValue(
|
||||
10 * (2**20), # 10MB
|
||||
environ_name="ATTACHMENT_MAX_SIZE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
MALWARE_DETECTION = {
|
||||
"BACKEND": values.Value(
|
||||
"lasuite.malware_detection.backends.dummy.DummyBackend",
|
||||
environ_name="MALWARE_DETECTION_BACKEND",
|
||||
environ_prefix=None,
|
||||
),
|
||||
"PARAMETERS": values.DictValue(
|
||||
default={
|
||||
"callback_path": "core.file_upload.malware_detection.malware_detection_callback",
|
||||
},
|
||||
environ_name="MALWARE_DETECTION_PARAMETERS",
|
||||
environ_prefix=None,
|
||||
),
|
||||
}
|
||||
|
||||
# Internationalization
|
||||
# https://docs.djangoproject.com/en/3.1/topics/i18n/
|
||||
|
||||
@@ -267,7 +141,7 @@ class Base(BraveSettings, Configuration):
|
||||
("en-us", "English"),
|
||||
("fr-fr", "Français"),
|
||||
# ("de-de", "Deutsch"),
|
||||
("nl-nl", "Nederlands"),
|
||||
# ("nl-nl", "Nederlands"),
|
||||
# ("es-es", "Español"),
|
||||
)
|
||||
)
|
||||
@@ -313,6 +187,7 @@ class Base(BraveSettings, Configuration):
|
||||
"django.middleware.csrf.CsrfViewMiddleware",
|
||||
"django.contrib.auth.middleware.AuthenticationMiddleware",
|
||||
"posthog.integrations.django.PosthogContextMiddleware",
|
||||
"core.middleware.PostHogMiddleware",
|
||||
"django.contrib.messages.middleware.MessageMiddleware",
|
||||
"dockerflow.django.middleware.DockerflowMiddleware",
|
||||
]
|
||||
@@ -375,16 +250,6 @@ class Base(BraveSettings, Configuration):
|
||||
"DEFAULT_VERSIONING_CLASS": "rest_framework.versioning.URLPathVersioning",
|
||||
"DEFAULT_SCHEMA_CLASS": "drf_spectacular.openapi.AutoSchema",
|
||||
"DEFAULT_THROTTLE_RATES": {
|
||||
"attachment_upload": values.Value(
|
||||
default="60/minute",
|
||||
environ_name="API_ATTACHMENT_UPLOAD_THROTTLE_RATE",
|
||||
environ_prefix=None,
|
||||
),
|
||||
"attachment_auth": values.Value(
|
||||
default="60/minute",
|
||||
environ_name="API_ATTACHMENT_AUTH_THROTTLE_RATE",
|
||||
environ_prefix=None,
|
||||
),
|
||||
"user_list_sustained": values.Value(
|
||||
default="180/hour",
|
||||
environ_name="API_USERS_LIST_THROTTLE_RATE_SUSTAINED",
|
||||
@@ -482,11 +347,7 @@ class Base(BraveSettings, Configuration):
|
||||
THUMBNAIL_ALIASES = {}
|
||||
|
||||
# Session
|
||||
SESSION_ENGINE = values.Value(
|
||||
"django.contrib.sessions.backends.cache",
|
||||
environ_name="SESSION_ENGINE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
SESSION_ENGINE = "django.contrib.sessions.backends.cache"
|
||||
SESSION_CACHE_ALIAS = "default"
|
||||
SESSION_COOKIE_AGE = values.PositiveIntegerValue(
|
||||
default=60 * 60 * 12, environ_name="SESSION_COOKIE_AGE", environ_prefix=None
|
||||
@@ -506,7 +367,6 @@ class Base(BraveSettings, Configuration):
|
||||
environ_name="OIDC_RP_CLIENT_SECRET",
|
||||
environ_prefix=None,
|
||||
)
|
||||
OIDC_OP_URL = values.Value(None, environ_name="OIDC_OP_URL", environ_prefix=None)
|
||||
OIDC_OP_JWKS_ENDPOINT = values.Value(environ_name="OIDC_OP_JWKS_ENDPOINT", environ_prefix=None)
|
||||
OIDC_OP_AUTHORIZATION_ENDPOINT = values.Value(
|
||||
environ_name="OIDC_OP_AUTHORIZATION_ENDPOINT", environ_prefix=None
|
||||
@@ -606,22 +466,6 @@ class Base(BraveSettings, Configuration):
|
||||
default=False, environ_name="ACTIVATION_REQUIRED", environ_prefix=None
|
||||
)
|
||||
|
||||
BREVO_API_KEY = values.Value(
|
||||
default=None,
|
||||
environ_name="BREVO_API_KEY",
|
||||
environ_prefix=None,
|
||||
)
|
||||
BREVO_FOLLOWUP_LIST_ID = values.Value(
|
||||
default=None,
|
||||
environ_name="BREVO_FOLLOWUP_LIST_ID",
|
||||
environ_prefix=None,
|
||||
)
|
||||
BREVO_WAITING_LIST_ID = values.Value(
|
||||
default=None,
|
||||
environ_name="BREVO_WAITING_LIST_ID",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
# AI service
|
||||
_llm_configuration_file_path = values.Value(
|
||||
os.path.join(BASE_DIR, "conversations/configuration/llm/default.json"),
|
||||
@@ -631,6 +475,9 @@ class Base(BraveSettings, Configuration):
|
||||
LLM_DEFAULT_MODEL_HRID = values.Value(
|
||||
"default-model", environ_name="LLM_DEFAULT_MODEL_HRID", environ_prefix=None
|
||||
)
|
||||
LLM_ROUTING_MODEL_HRID = values.Value(
|
||||
"default-routing-model", environ_name="LLM_ROUTING_MODEL_HRID", environ_prefix=None
|
||||
)
|
||||
LLM_SUMMARIZATION_MODEL_HRID = values.Value(
|
||||
"default-summarization-model",
|
||||
environ_name="LLM_SUMMARIZATION_MODEL_HRID",
|
||||
@@ -717,11 +564,6 @@ class Base(BraveSettings, Configuration):
|
||||
environ_name="RAG_DOCUMENT_SEARCH_BACKEND",
|
||||
environ_prefix=None,
|
||||
)
|
||||
SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = values.DictValue(
|
||||
default={},
|
||||
environ_name="SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
# Web search
|
||||
RAG_WEB_SEARCH_PROMPT_UPDATE = values.Value(
|
||||
@@ -791,21 +633,6 @@ USER QUESTION:
|
||||
environ_name="SUMMARIZATION_SYSTEM_PROMPT",
|
||||
environ_prefix=None,
|
||||
)
|
||||
SUMMARIZATION_CHUNK_SIZE = values.PositiveIntegerValue(
|
||||
default=20_000, # Approx 20k words per chunk
|
||||
environ_name="SUMMARIZATION_CHUNK_SIZE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
SUMMARIZATION_OVERLAP_SIZE = values.FloatValue(
|
||||
default=0.05, # 5% overlap
|
||||
environ_name="SUMMARIZATION_OVERLAP_SIZE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
SUMMARIZATION_CONCURRENT_REQUESTS = values.PositiveIntegerValue(
|
||||
default=5,
|
||||
environ_name="SUMMARIZATION_CONCURRENT_REQUESTS",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
# Tavily API
|
||||
TAVILY_API_KEY = values.Value(
|
||||
@@ -911,15 +738,6 @@ USER QUESTION:
|
||||
LANGFUSE_MEDIA_UPLOAD_ENABLED = values.BooleanValue(
|
||||
default=False, environ_name="LANGFUSE_MEDIA_UPLOAD_ENABLED", environ_prefix=None
|
||||
)
|
||||
AUTO_TITLE_AFTER_USER_MESSAGES = values.PositiveIntegerValue(
|
||||
3, environ_name="AUTO_TITLE_AFTER_USER_MESSAGES", environ_prefix=None
|
||||
)
|
||||
# WARNING: Testing purpose only. Do not use in production.
|
||||
WARNING_MOCK_CONVERSATION_AGENT = values.BooleanValue(
|
||||
default=False,
|
||||
environ_name="WARNING_MOCK_CONVERSATION_AGENT",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
@property
|
||||
@@ -1062,17 +880,6 @@ class Build(Base):
|
||||
"""
|
||||
|
||||
SECRET_KEY = values.Value("DummyKey")
|
||||
STORAGES = {
|
||||
"default": {
|
||||
"BACKEND": "django.core.files.storage.FileSystemStorage",
|
||||
},
|
||||
"staticfiles": {
|
||||
"BACKEND": values.Value(
|
||||
"whitenoise.storage.CompressedManifestStaticFilesStorage",
|
||||
environ_name="STORAGES_STATICFILES_BACKEND",
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class Development(Base):
|
||||
|
||||
@@ -13,6 +13,5 @@ from configurations.wsgi import get_wsgi_application
|
||||
|
||||
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "conversations.settings")
|
||||
os.environ.setdefault("DJANGO_CONFIGURATION", "Development")
|
||||
os.environ.setdefault("PYTHON_SERVER_MODE", "sync")
|
||||
|
||||
application = get_wsgi_application()
|
||||
|
||||
@@ -17,9 +17,8 @@ class UserSerializer(serializers.ModelSerializer):
|
||||
"full_name",
|
||||
"short_name",
|
||||
"language",
|
||||
"sub",
|
||||
]
|
||||
read_only_fields = ["id", "email", "full_name", "short_name", "sub"]
|
||||
read_only_fields = ["id", "email", "full_name", "short_name"]
|
||||
|
||||
|
||||
class UserLightSerializer(UserSerializer):
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Authentication Backends for the Conversations core app."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import SuspiciousOperation
|
||||
@@ -9,11 +10,23 @@ from lasuite.oidc_login.backends import (
|
||||
OIDCAuthenticationBackend as LaSuiteOIDCAuthenticationBackend,
|
||||
)
|
||||
|
||||
from core.brevo import add_user_to_brevo_list
|
||||
from core.models import DuplicateEmailError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Settings renamed warnings
|
||||
if os.environ.get("USER_OIDC_FIELDS_TO_FULLNAME"):
|
||||
logger.warning(
|
||||
"USER_OIDC_FIELDS_TO_FULLNAME has been renamed "
|
||||
"to OIDC_USERINFO_FULLNAME_FIELDS please update your settings."
|
||||
)
|
||||
|
||||
if os.environ.get("USER_OIDC_FIELD_TO_SHORTNAME"):
|
||||
logger.warning(
|
||||
"USER_OIDC_FIELD_TO_SHORTNAME has been renamed "
|
||||
"to OIDC_USERINFO_SHORTNAME_FIELD please update your settings."
|
||||
)
|
||||
|
||||
|
||||
class OIDCAuthenticationBackend(LaSuiteOIDCAuthenticationBackend):
|
||||
"""Custom OpenID Connect (OIDC) Authentication Backend.
|
||||
@@ -57,12 +70,3 @@ class OIDCAuthenticationBackend(LaSuiteOIDCAuthenticationBackend):
|
||||
return super().create_user(
|
||||
claims | {"allow_conversation_analytics": settings.DEFAULT_ALLOW_CONVERSATION_ANALYTICS}
|
||||
)
|
||||
|
||||
def authenticate(self, request, **kwargs):
|
||||
"""Authenticate user and add they to Brevo list if activation not required."""
|
||||
user = super().authenticate(request, **kwargs)
|
||||
|
||||
if user and not settings.ACTIVATION_REQUIRED and settings.BREVO_FOLLOWUP_LIST_ID:
|
||||
add_user_to_brevo_list([user.email], settings.BREVO_FOLLOWUP_LIST_ID)
|
||||
|
||||
return user
|
||||
|
||||
@@ -1,135 +0,0 @@
|
||||
"""Functions for interacting with Brevo API to manage contacts in a waiting list."""
|
||||
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def create_contact_in_brevo(email: str) -> bool:
|
||||
"""
|
||||
Create a contact in Brevo.
|
||||
|
||||
Args:
|
||||
email (str): The email address of the user.
|
||||
|
||||
"""
|
||||
api_key = settings.BREVO_API_KEY
|
||||
if not api_key:
|
||||
logger.info("Brevo API key not configured: skipping creating contact")
|
||||
return False
|
||||
|
||||
url = "https://api.brevo.com/v3/contacts"
|
||||
headers = {
|
||||
"accept": "application/json",
|
||||
"api-key": api_key,
|
||||
"content-type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"email": email,
|
||||
"updateEnabled": True,
|
||||
}
|
||||
try:
|
||||
response = requests.post(url, json=payload, headers=headers, timeout=5)
|
||||
except requests.RequestException as e:
|
||||
logger.exception(e)
|
||||
return False
|
||||
|
||||
if not response.ok:
|
||||
logger.error(
|
||||
"Error creating contact in Brevo %s: (%s) %s",
|
||||
email,
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def add_user_to_brevo_list(emails: List[str], list_id: Optional[str]) -> None:
|
||||
"""
|
||||
Add email list to a Brevo list.
|
||||
|
||||
Args:
|
||||
emails (List[str]): The email address(es) of the user(s).
|
||||
list_id (str): The Brevo waiting list ID, can be None if not configured.
|
||||
|
||||
"""
|
||||
api_key = settings.BREVO_API_KEY
|
||||
if not api_key or not list_id:
|
||||
logger.info("Brevo API key or list ID not configured: skipping adding contact")
|
||||
return
|
||||
|
||||
for email in emails:
|
||||
# Ensure the contact exists before adding to the list
|
||||
# `emails` contains several emails only when used from the admin interface bulk action
|
||||
if not create_contact_in_brevo(email):
|
||||
logger.error("Failed to create contact %s in Brevo, skipping adding to list", email)
|
||||
return
|
||||
|
||||
url = f"https://api.brevo.com/v3/contacts/lists/{list_id}/contacts/add"
|
||||
headers = {
|
||||
"accept": "application/json",
|
||||
"api-key": api_key,
|
||||
"content-type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"emails": emails,
|
||||
}
|
||||
try:
|
||||
response = requests.post(url, json=payload, headers=headers, timeout=5)
|
||||
except requests.RequestException as e:
|
||||
logger.exception(e)
|
||||
return
|
||||
|
||||
if response.status_code != 201:
|
||||
logger.error(
|
||||
"Error adding contacts to Brevo (%s) %s: (%s) %s",
|
||||
list_id,
|
||||
emails,
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
|
||||
|
||||
def remove_user_from_brevo_list(emails: List[str], list_id: Optional[str]) -> None:
|
||||
"""
|
||||
Remove email list from a Brevo list.
|
||||
|
||||
Args:
|
||||
emails (List[str]): The email address(es) of the user(s).
|
||||
list_id (str): The Brevo waiting list ID, can be None if not configured.
|
||||
|
||||
"""
|
||||
api_key = settings.BREVO_API_KEY
|
||||
if not api_key or not list_id:
|
||||
logger.info("Brevo API key or list ID not configured: skipping removing contact")
|
||||
return
|
||||
|
||||
url = f"https://api.brevo.com/v3/contacts/lists/{list_id}/contacts/remove"
|
||||
headers = {
|
||||
"accept": "application/json",
|
||||
"api-key": api_key,
|
||||
"content-type": "application/json",
|
||||
}
|
||||
payload = {
|
||||
"emails": emails,
|
||||
}
|
||||
try:
|
||||
response = requests.post(url, json=payload, headers=headers, timeout=5)
|
||||
except requests.RequestException as e:
|
||||
logger.exception(e)
|
||||
return
|
||||
if response.status_code != 201:
|
||||
logger.error(
|
||||
"Error removing contacts from Brevo (%s) %s: (%s) %s",
|
||||
list_id,
|
||||
emails,
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
@@ -1,8 +1,22 @@
|
||||
"""Core application enums declaration."""
|
||||
|
||||
from django.conf import global_settings
|
||||
import re
|
||||
|
||||
from django.conf import global_settings, settings
|
||||
from django.utils.translation import gettext_lazy as _
|
||||
|
||||
ATTACHMENTS_FOLDER = "attachments"
|
||||
UUID_REGEX = r"[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}"
|
||||
FILE_EXT_REGEX = r"\.[a-zA-Z0-9]{1,10}"
|
||||
MEDIA_STORAGE_URL_PATTERN = re.compile(
|
||||
f"{settings.MEDIA_URL:s}(?P<pk>{UUID_REGEX:s})/"
|
||||
f"(?P<attachment>{ATTACHMENTS_FOLDER:s}/{UUID_REGEX:s}(?:-unsafe)?{FILE_EXT_REGEX:s})$"
|
||||
)
|
||||
MEDIA_STORAGE_URL_EXTRACT = re.compile(
|
||||
f"{settings.MEDIA_URL:s}({UUID_REGEX}/{ATTACHMENTS_FOLDER}/{UUID_REGEX}{FILE_EXT_REGEX})"
|
||||
)
|
||||
|
||||
|
||||
# In Django's code base, `LANGUAGES` is set by default with all supported languages.
|
||||
# We can use it for the choice of languages which should not be limited to the few languages
|
||||
# active in the app.
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
from enum import StrEnum
|
||||
|
||||
from django.conf import settings
|
||||
from django.utils.text import slugify
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
@@ -44,9 +43,3 @@ class FeatureFlags(BaseModel):
|
||||
# features
|
||||
web_search: FeatureToggle = FeatureToggle.DISABLED
|
||||
document_upload: FeatureToggle = FeatureToggle.DISABLED
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
"""Dynamically get specific RAG document search tool feature flags from settings."""
|
||||
if config := settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS.get(name):
|
||||
return FeatureToggle[config.get("feature_flag_value", "DISABLED").upper()]
|
||||
return super().__getattr__(name)
|
||||
|
||||
@@ -9,6 +9,7 @@ User = get_user_model()
|
||||
|
||||
try:
|
||||
import posthog
|
||||
from posthog.contexts import get_tags
|
||||
except ImportError:
|
||||
posthog = None
|
||||
|
||||
@@ -38,7 +39,8 @@ def is_feature_enabled(
|
||||
if posthog is not None:
|
||||
return posthog.feature_enabled(
|
||||
frontend_feature_name(feature_name),
|
||||
str(user.pk), # same as set by the frontend
|
||||
user.email,
|
||||
person_properties={"$host": get_tags().get("$host")},
|
||||
)
|
||||
|
||||
# No feature flag manager
|
||||
|
||||
@@ -1,22 +0,0 @@
|
||||
"""File upload management enums declaration."""
|
||||
|
||||
from enum import StrEnum
|
||||
|
||||
UUID_REGEX = r"[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}"
|
||||
FILE_EXT_REGEX = r"\.[a-zA-Z0-9]{1,10}"
|
||||
|
||||
|
||||
class AttachmentStatus(StrEnum):
|
||||
"""Defines the possible statuses for an attachment."""
|
||||
|
||||
PENDING = "pending"
|
||||
UPLOADING = "uploading"
|
||||
ANALYZING = "analyzing"
|
||||
FILE_TOO_LARGE_TO_ANALYZE = "file_too_large_to_analyze"
|
||||
SUSPICIOUS = "suspicious"
|
||||
READY = "ready"
|
||||
|
||||
@classmethod
|
||||
def choices(cls):
|
||||
"""Return a list of tuples for each enum member."""
|
||||
return [(member.value, member.name) for member in cls]
|
||||
@@ -1,5 +0,0 @@
|
||||
"""Defines exceptions for file upload operations."""
|
||||
|
||||
|
||||
class HolderDoesNotExist(Exception):
|
||||
"""Raised when the holder of an attachment does not exist."""
|
||||
@@ -1,28 +0,0 @@
|
||||
"""Malware detection callbacks"""
|
||||
|
||||
import logging
|
||||
|
||||
from django.utils.module_loading import import_string
|
||||
|
||||
from lasuite.malware_detection.enums import ReportStatus
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def malware_detection_callback( # noqa: PLR0913 # pylint: disable=too-many-arguments
|
||||
file_path, status, error_info, *, safe_callback, unknown_callback, unsafe_callback, **kwargs
|
||||
):
|
||||
"""Malware detection callback"""
|
||||
|
||||
if status == ReportStatus.SAFE:
|
||||
_safe_callback = import_string(safe_callback)
|
||||
_safe_callback(file_path, **kwargs)
|
||||
return
|
||||
|
||||
if status == ReportStatus.UNKNOWN:
|
||||
_unknown_callback = import_string(unknown_callback)
|
||||
if _unknown_callback(file_path, error_info, **kwargs):
|
||||
return
|
||||
|
||||
_unsafe_callback = import_string(unsafe_callback)
|
||||
_unsafe_callback(file_path, error_info, **kwargs)
|
||||
@@ -1,307 +0,0 @@
|
||||
"""Mixin to add attachment upload and access control to a viewset."""
|
||||
|
||||
import logging
|
||||
import re
|
||||
import uuid
|
||||
from urllib.parse import quote, urlencode
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
from django.utils.text import get_valid_filename
|
||||
|
||||
from botocore.exceptions import ClientError
|
||||
from lasuite.malware_detection import malware_detection
|
||||
from rest_framework import decorators, exceptions, status
|
||||
from rest_framework.response import Response
|
||||
from rest_framework.reverse import reverse
|
||||
from rest_framework.throttling import UserRateThrottle
|
||||
|
||||
from . import enums, utils
|
||||
from .enums import FILE_EXT_REGEX, UUID_REGEX
|
||||
from .exceptions import HolderDoesNotExist
|
||||
from .serializers import FileUploadSerializer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AttachmentUploadThrottle(UserRateThrottle):
|
||||
"""Throttle for the attachment upload endpoint."""
|
||||
|
||||
scope = "attachment_upload"
|
||||
|
||||
|
||||
class AttachmentAuthThrottle(UserRateThrottle):
|
||||
"""Throttle for the attachment upload endpoint."""
|
||||
|
||||
scope = "attachment_auth"
|
||||
|
||||
|
||||
class AttachmentMixin:
|
||||
"""
|
||||
Mixin to add attachment upload and access control to a viewset.
|
||||
|
||||
The viewset must be based on `GenericViewSet` and define the following methods because
|
||||
they are highly model-specific and depends on how attachments are stored (a key in
|
||||
an ArrayField, a related model, etc.):
|
||||
- `malware_detection_kwargs`
|
||||
- `store_attachment`
|
||||
- `get_holder_from_key`
|
||||
- `_check_attachment_present`
|
||||
"""
|
||||
|
||||
ATTACHMENTS_FOLDER = "attachments"
|
||||
MEDIA_STORAGE_URL_PATTERN = re.compile(
|
||||
f"{settings.MEDIA_URL:s}(?P<pk>{UUID_REGEX:s})/"
|
||||
f"(?P<attachment>{ATTACHMENTS_FOLDER:s}/{UUID_REGEX:s}(?:-unsafe)?{FILE_EXT_REGEX:s})$"
|
||||
)
|
||||
|
||||
def malware_detection_kwargs(self, holder) -> dict:
|
||||
"""
|
||||
Extra arguments to pass to the malware detection backend.
|
||||
The result will be passed to the `analyse_file` callback method of the backend.
|
||||
|
||||
Should return a dictionary like `{"document_id": str(holder.pk)}`
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def store_attachment(self, holder, key, serializer) -> None:
|
||||
"""
|
||||
Store the attachment key and save it: this provides the permission to access it.
|
||||
|
||||
Used by the attachment_upload endpoint.
|
||||
Can be like:
|
||||
|
||||
```
|
||||
holder.attachments.add(key)
|
||||
holder.save()
|
||||
```
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_holder_from_key(self, url_params, key):
|
||||
"""
|
||||
Get the holder object from the attachment key, to check the user has
|
||||
access to the attachment holder (so they have access to the attachment).
|
||||
|
||||
Used by the media_auth endpoint.
|
||||
Can be like:
|
||||
|
||||
```
|
||||
return self.queryset.get(attachments__contains=[key])
|
||||
```
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def _check_attachment_present(self, holder, key) -> bool:
|
||||
"""
|
||||
Check if the attachment key is present in the holder's attachments.
|
||||
|
||||
Used by the media_check endpoint.
|
||||
Can be like:
|
||||
|
||||
```
|
||||
return key in holder.attachments
|
||||
```
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_object_key_base(self, holder):
|
||||
"""Key base of the location where the attachment is stored in object storage."""
|
||||
return str(holder.pk)
|
||||
|
||||
def get_media_check_url(self, holder):
|
||||
"""Get the URL to check the status of an attachment."""
|
||||
return reverse(
|
||||
f"{self.basename}-media-check",
|
||||
kwargs={"pk": holder.pk},
|
||||
)
|
||||
|
||||
def check_attachment_holder_permission(self, user, url_params, key):
|
||||
"""
|
||||
Check if the user has permission to access the holder of the attachment.
|
||||
|
||||
Raises PermissionDenied if the user does not have permission.
|
||||
"""
|
||||
try:
|
||||
holder = self.get_holder_from_key(url_params, key)
|
||||
except HolderDoesNotExist as err:
|
||||
logger.debug("Attachment holder not found for key '%s': %s", key, err)
|
||||
# We raise PermissionDenied instead of Http404 to avoid leaking information
|
||||
# about the existence of the attachment.
|
||||
raise exceptions.PermissionDenied() from err
|
||||
|
||||
self.check_object_permissions(self.request, holder)
|
||||
# for now, only the owner can access the attachment
|
||||
if holder.owner_id != user.pk:
|
||||
raise exceptions.PermissionDenied()
|
||||
|
||||
def get_key_from_url_params(self, url_params) -> str:
|
||||
"""
|
||||
Get the attachment key from URL parameters.
|
||||
|
||||
The key is composed of the holder's pk and the attachment filename.
|
||||
"""
|
||||
return f"{url_params['pk']:s}/{url_params['attachment']:s}"
|
||||
|
||||
def check_attachment_is_ready(self, key) -> bool:
|
||||
"""
|
||||
Check if the attachment is ready to be served.
|
||||
|
||||
An attachment is ready if its metadata `status` is `READY`.
|
||||
"""
|
||||
s3_client = default_storage.connection.meta.client
|
||||
bucket_name = default_storage.bucket_name
|
||||
try:
|
||||
head_resp = s3_client.head_object(Bucket=bucket_name, Key=key)
|
||||
except ClientError as err:
|
||||
logger.error("Client Error fetching file %s metadata: %s", key, err)
|
||||
return False
|
||||
|
||||
metadata = head_resp.get("Metadata", {})
|
||||
# In order to be compatible with existing upload without `status` metadata,
|
||||
# we consider them as ready.
|
||||
return metadata.get("status", enums.AttachmentStatus.READY) == enums.AttachmentStatus.READY
|
||||
|
||||
@decorators.action(detail=True, methods=["post"], url_path="attachment-upload")
|
||||
@decorators.throttle_classes([AttachmentUploadThrottle])
|
||||
def attachment_upload(self, request, *args, **kwargs):
|
||||
"""Upload a file related to a given document"""
|
||||
# Check permissions first
|
||||
holder = self.get_object()
|
||||
|
||||
# Validate metadata in payload
|
||||
serializer = FileUploadSerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
|
||||
# Generate a generic yet unique filename to store the image in object storage
|
||||
file_id = uuid.uuid4()
|
||||
ext = serializer.validated_data["expected_extension"]
|
||||
|
||||
# Prepare metadata for storage
|
||||
extra_args = {
|
||||
"Metadata": {
|
||||
"owner": str(request.user.id),
|
||||
"status": enums.AttachmentStatus.ANALYZING,
|
||||
},
|
||||
"ContentType": serializer.validated_data["content_type"],
|
||||
}
|
||||
file_unsafe = ""
|
||||
if serializer.validated_data["is_unsafe"]:
|
||||
extra_args["Metadata"]["is_unsafe"] = "true"
|
||||
file_unsafe = "-unsafe"
|
||||
|
||||
holder_key_base = self.get_object_key_base(holder)
|
||||
key = f"{holder_key_base}/{self.ATTACHMENTS_FOLDER:s}/{file_id!s}{file_unsafe}.{ext:s}"
|
||||
|
||||
raw_name = serializer.validated_data["file_name"]
|
||||
# Strip CR/LF and normalize to a filesystem-safe basename
|
||||
safe_name = get_valid_filename(raw_name.replace("\r", "").replace("\n", "")) or "file"
|
||||
ascii_name = safe_name.encode("ascii", "ignore").decode("ascii") or "file"
|
||||
# RFC 5987 filename* for non-ASCII
|
||||
disp_filename = f'filename="{ascii_name}"'
|
||||
disp_filename_star = f"filename*=UTF-8''{quote(safe_name)}"
|
||||
if (
|
||||
not serializer.validated_data["content_type"].startswith("image/")
|
||||
or serializer.validated_data["is_unsafe"]
|
||||
):
|
||||
extra_args.update(
|
||||
{"ContentDisposition": f"attachment; {disp_filename}; {disp_filename_star}"}
|
||||
)
|
||||
else:
|
||||
extra_args.update(
|
||||
{"ContentDisposition": f"inline; {disp_filename}; {disp_filename_star}"}
|
||||
)
|
||||
|
||||
file = serializer.validated_data["file"]
|
||||
default_storage.connection.meta.client.upload_fileobj(
|
||||
file, default_storage.bucket_name, key, ExtraArgs=extra_args
|
||||
)
|
||||
|
||||
self.store_attachment(holder, key, serializer)
|
||||
|
||||
malware_detection.analyse_file(key, **self.malware_detection_kwargs(holder))
|
||||
|
||||
url = self.get_media_check_url(holder)
|
||||
parameters = urlencode({"key": key})
|
||||
|
||||
return Response(
|
||||
{
|
||||
"file": f"{url:s}?{parameters:s}",
|
||||
},
|
||||
status=status.HTTP_201_CREATED,
|
||||
)
|
||||
|
||||
@decorators.action(detail=False, methods=["get"], url_path="media-auth")
|
||||
@decorators.throttle_classes([AttachmentAuthThrottle])
|
||||
def media_auth(self, request, *args, **kwargs):
|
||||
"""
|
||||
This view is used by an Nginx subrequest to control access to a document's
|
||||
attachment file.
|
||||
|
||||
When we let the request go through, we compute authorization headers that will be added to
|
||||
the request going through thanks to the nginx.ingress.kubernetes.io/auth-response-headers
|
||||
annotation. The request will then be proxied to the object storage backend who will
|
||||
respond with the file after checking the signature included in headers.
|
||||
"""
|
||||
parsed_url = utils.auth_get_original_url(request)
|
||||
url_params = utils.auth_get_url_params(self.MEDIA_STORAGE_URL_PATTERN, parsed_url.path)
|
||||
key = self.get_key_from_url_params(url_params)
|
||||
|
||||
user = request.user
|
||||
|
||||
# Look for a holder to which the user has access and that includes this attachment
|
||||
# Might raise PermissionDenied
|
||||
self.check_attachment_holder_permission(user, url_params, key)
|
||||
|
||||
# Check if the attachment is ready
|
||||
if not self.check_attachment_is_ready(key):
|
||||
raise exceptions.PermissionDenied()
|
||||
|
||||
# Generate S3 authorization headers using the extracted URL parameters
|
||||
request = utils.generate_s3_authorization_headers(key)
|
||||
|
||||
return Response("authorized", headers=request.headers, status=status.HTTP_200_OK)
|
||||
|
||||
@decorators.action(detail=True, methods=["get"], url_path="media-check")
|
||||
def media_check(self, request, *args, **kwargs):
|
||||
"""
|
||||
Check if the media is ready to be served.
|
||||
"""
|
||||
holder = self.get_object()
|
||||
|
||||
key = request.query_params.get("key")
|
||||
if not key:
|
||||
return Response(
|
||||
{"detail": "Missing 'key' query parameter"},
|
||||
status=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
if not self._check_attachment_present(holder, key):
|
||||
return Response(
|
||||
{"detail": "Attachment missing"},
|
||||
status=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
|
||||
# Check if the attachment is ready
|
||||
s3_client = default_storage.connection.meta.client
|
||||
bucket_name = default_storage.bucket_name
|
||||
try:
|
||||
head_resp = s3_client.head_object(Bucket=bucket_name, Key=key)
|
||||
except ClientError as err:
|
||||
logger.error("Client Error fetching file %s metadata: %s", key, err)
|
||||
return Response(
|
||||
{"detail": "Media not found"},
|
||||
status=status.HTTP_404_NOT_FOUND,
|
||||
)
|
||||
metadata = head_resp.get("Metadata", {})
|
||||
|
||||
body = {
|
||||
"status": metadata.get("status", enums.AttachmentStatus.ANALYZING),
|
||||
}
|
||||
if metadata.get("status") == enums.AttachmentStatus.READY:
|
||||
body = {
|
||||
"status": enums.AttachmentStatus.READY,
|
||||
"file": f"{settings.MEDIA_URL:s}{key:s}",
|
||||
}
|
||||
|
||||
return Response(body, status=status.HTTP_200_OK)
|
||||
@@ -1,62 +0,0 @@
|
||||
"""Serializers for file upload API."""
|
||||
|
||||
import mimetypes
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
import magic
|
||||
from rest_framework import serializers
|
||||
|
||||
|
||||
class FileUploadSerializer(serializers.Serializer): # pylint: disable=abstract-method
|
||||
"""Receive file upload requests."""
|
||||
|
||||
file = serializers.FileField()
|
||||
|
||||
def validate_file(self, file):
|
||||
"""Add file size and type constraints as defined in settings."""
|
||||
# Validate file size
|
||||
if file.size > settings.ATTACHMENT_MAX_SIZE:
|
||||
max_size = settings.ATTACHMENT_MAX_SIZE // (1024 * 1024)
|
||||
raise serializers.ValidationError(
|
||||
f"File size exceeds the maximum limit of {max_size:d} MB."
|
||||
)
|
||||
|
||||
extension = file.name.rpartition(".")[-1] if "." in file.name else None
|
||||
|
||||
# Read the first few bytes to determine the MIME type accurately
|
||||
mime = magic.Magic(mime=True)
|
||||
magic_mime_type = mime.from_buffer(file.read(1024))
|
||||
file.seek(0) # Reset file pointer to the beginning after reading
|
||||
self.context["is_unsafe"] = False
|
||||
if settings.ATTACHMENT_CHECK_UNSAFE_MIME_TYPES_ENABLED:
|
||||
self.context["is_unsafe"] = magic_mime_type in settings.ATTACHMENT_UNSAFE_MIME_TYPES
|
||||
|
||||
extension_mime_type, _ = mimetypes.guess_type(file.name)
|
||||
|
||||
# Try guessing a coherent extension from the mimetype
|
||||
if extension_mime_type != magic_mime_type:
|
||||
self.context["is_unsafe"] = True
|
||||
|
||||
guessed_ext = mimetypes.guess_extension(magic_mime_type)
|
||||
# Missing extensions or extensions longer than 5 characters (it's as long as an extension
|
||||
# can be) are replaced by the extension we eventually guessed from mimetype.
|
||||
if (extension is None or len(extension) > 5) and guessed_ext:
|
||||
extension = guessed_ext[1:]
|
||||
|
||||
if extension is None:
|
||||
raise serializers.ValidationError("Could not determine file extension.")
|
||||
|
||||
self.context["expected_extension"] = extension
|
||||
self.context["content_type"] = magic_mime_type
|
||||
self.context["file_name"] = file.name
|
||||
|
||||
return file
|
||||
|
||||
def validate(self, attrs):
|
||||
"""Override validate to add the computed extension to validated_data."""
|
||||
attrs["expected_extension"] = self.context["expected_extension"]
|
||||
attrs["is_unsafe"] = self.context["is_unsafe"]
|
||||
attrs["content_type"] = self.context["content_type"]
|
||||
attrs["file_name"] = self.context["file_name"]
|
||||
return attrs
|
||||
@@ -1,150 +0,0 @@
|
||||
"""Util to generate S3 authorization headers for object storage access control"""
|
||||
|
||||
import logging
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
import boto3
|
||||
import botocore
|
||||
from rest_framework import exceptions
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def auth_get_original_url(request):
|
||||
"""
|
||||
Extracts and parses the original URL from the "HTTP_X_ORIGINAL_URL" header.
|
||||
Raises PermissionDenied if the header is missing.
|
||||
|
||||
The original url is passed by nginx in the "HTTP_X_ORIGINAL_URL" header.
|
||||
See corresponding ingress configuration in Helm chart and read about the
|
||||
nginx.ingress.kubernetes.io/auth-url annotation to understand how the Nginx ingress
|
||||
is configured to do this.
|
||||
|
||||
Based on the original url and the logged in user, we must decide if we authorize Nginx
|
||||
to let this request go through (by returning a 200 code) or if we block it (by returning
|
||||
a 403 error). Note that we return 403 errors without any further details for security
|
||||
reasons.
|
||||
"""
|
||||
# Extract the original URL from the request header
|
||||
original_url = request.META.get("HTTP_X_ORIGINAL_URL")
|
||||
if not original_url:
|
||||
logger.debug("Missing HTTP_X_ORIGINAL_URL header in subrequest")
|
||||
raise exceptions.PermissionDenied()
|
||||
|
||||
logger.debug("Original url: '%s'", original_url)
|
||||
return urlparse(original_url)
|
||||
|
||||
|
||||
def auth_get_url_params(pattern, fragment):
|
||||
"""
|
||||
Extracts URL parameters from the given fragment using the specified regex pattern.
|
||||
Raises PermissionDenied if parameters cannot be extracted.
|
||||
"""
|
||||
match = pattern.search(fragment)
|
||||
try:
|
||||
return match.groupdict()
|
||||
except (ValueError, AttributeError) as exc:
|
||||
logger.debug("Failed to extract parameters from subrequest URL: %s", exc)
|
||||
raise exceptions.PermissionDenied() from exc
|
||||
|
||||
|
||||
def generate_s3_authorization_headers(key):
|
||||
"""
|
||||
Generate authorization headers for an s3 object.
|
||||
These headers can be used as an alternative to signed urls with many benefits:
|
||||
- the urls of our files never expire and can be stored in our documents' content
|
||||
- we don't leak authorized urls that could be shared (file access can only be done
|
||||
with cookies)
|
||||
- access control is truly realtime
|
||||
- the object storage service does not need to be exposed on internet
|
||||
"""
|
||||
url = default_storage.unsigned_connection.meta.client.generate_presigned_url(
|
||||
"get_object",
|
||||
ExpiresIn=0,
|
||||
Params={"Bucket": default_storage.bucket_name, "Key": key},
|
||||
)
|
||||
request = botocore.awsrequest.AWSRequest(method="get", url=url)
|
||||
|
||||
s3_client = default_storage.connection.meta.client
|
||||
# pylint: disable=protected-access
|
||||
credentials = s3_client._request_signer._credentials # noqa: SLF001
|
||||
frozen_credentials = credentials.get_frozen_credentials()
|
||||
region = s3_client.meta.region_name
|
||||
auth = botocore.auth.S3SigV4Auth(frozen_credentials, "s3", region)
|
||||
auth.add_auth(request)
|
||||
|
||||
return request
|
||||
|
||||
|
||||
def _get_s3_client() -> botocore.client.BaseClient:
|
||||
"""
|
||||
Get the S3 client according to the settings.
|
||||
|
||||
If AWS_S3_DOMAIN_REPLACE is set, create a new S3 client with the specified endpoint_url.
|
||||
Otherwise, use the existing client from the default storage.
|
||||
"""
|
||||
# This settings should be used if the backend application and the frontend application
|
||||
# can't connect to the object storage with the same domain. This is the case in the
|
||||
# docker compose stack used in development. The frontend application will use localhost
|
||||
# to connect to the object storage while the backend application will use the object storage
|
||||
# service name declared in the docker compose stack.
|
||||
# This is needed because the domain name is used to compute the signature. So it can't be
|
||||
# changed dynamically by the frontend application.
|
||||
if settings.AWS_S3_DOMAIN_REPLACE:
|
||||
return boto3.client(
|
||||
"s3",
|
||||
aws_access_key_id=settings.AWS_S3_ACCESS_KEY_ID,
|
||||
aws_secret_access_key=settings.AWS_S3_SECRET_ACCESS_KEY,
|
||||
endpoint_url=settings.AWS_S3_DOMAIN_REPLACE,
|
||||
config=botocore.client.Config(
|
||||
region_name=settings.AWS_S3_REGION_NAME,
|
||||
signature_version=settings.AWS_S3_SIGNATURE_VERSION,
|
||||
),
|
||||
)
|
||||
|
||||
return default_storage.connection.meta.client
|
||||
|
||||
|
||||
def generate_upload_policy(key: str):
|
||||
"""
|
||||
Generate a S3 upload policy for a given key.
|
||||
|
||||
Args:
|
||||
key (str): The S3 object key where the file will be uploaded.
|
||||
"""
|
||||
|
||||
# Get the S3 client according to the settings
|
||||
s3_client = _get_s3_client()
|
||||
|
||||
# Generate the policy
|
||||
policy = s3_client.generate_presigned_url(
|
||||
ClientMethod="put_object",
|
||||
Params={"Bucket": default_storage.bucket_name, "Key": key, "ACL": "private"},
|
||||
ExpiresIn=settings.AWS_S3_UPLOAD_POLICY_EXPIRATION,
|
||||
)
|
||||
|
||||
return policy
|
||||
|
||||
|
||||
def generate_retrieve_policy(key: str):
|
||||
"""
|
||||
Generate a S3 retrieve policy for a given item.
|
||||
|
||||
Args:
|
||||
key (str): The S3 object key where the file is stored.
|
||||
"""
|
||||
|
||||
# Get the S3 client according to the settings
|
||||
s3_client = _get_s3_client()
|
||||
|
||||
# Generate the policy
|
||||
policy = s3_client.generate_presigned_url(
|
||||
ClientMethod="get_object",
|
||||
Params={"Bucket": default_storage.bucket_name, "Key": key},
|
||||
ExpiresIn=settings.AWS_S3_RETRIEVE_POLICY_EXPIRATION,
|
||||
)
|
||||
|
||||
return policy
|
||||
@@ -1,65 +0,0 @@
|
||||
"""Management command to get or set the CORS configuration of the storage bucket"""
|
||||
|
||||
import json
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
from django.core.management.base import BaseCommand
|
||||
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
|
||||
class Command(BaseCommand):
|
||||
"""
|
||||
Management command to get or set the CORS configuration of the storage bucket
|
||||
"""
|
||||
|
||||
help = "Get or set the CORS configuration of the storage bucket based on the Django settings"
|
||||
|
||||
def add_arguments(self, parser):
|
||||
"""Adds the command-line argument to the command"""
|
||||
parser.add_argument("--set", action="store_true", help="Set the CORS configuration")
|
||||
|
||||
def handle(self, *args, **options):
|
||||
"""Handles the command"""
|
||||
|
||||
s3_client = default_storage.connection.meta.client
|
||||
|
||||
if options["set"]:
|
||||
# Set CORS rules
|
||||
if "*" in settings.ALLOWED_HOSTS:
|
||||
allowed_origins = ["*"]
|
||||
elif len(settings.ALLOWED_HOSTS) > 0:
|
||||
allowed_origins = [f"https://{h}" for h in settings.ALLOWED_HOSTS]
|
||||
else:
|
||||
raise ValueError("DJANGO_ALLOWED_HOSTS is not set")
|
||||
|
||||
cors_config = {
|
||||
"CORSRules": [
|
||||
{
|
||||
"AllowedOrigins": allowed_origins,
|
||||
"AllowedHeaders": ["*"],
|
||||
"AllowedMethods": ["GET", "HEAD", "POST", "PUT", "DELETE"],
|
||||
"MaxAgeSeconds": 3000,
|
||||
"ExposeHeaders": ["Etag"],
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
s3_client.put_bucket_cors(
|
||||
Bucket=default_storage.bucket_name, CORSConfiguration=cors_config
|
||||
)
|
||||
self.stdout.write(self.style.SUCCESS("CORS configuration successfully updated."))
|
||||
else:
|
||||
# Get CORS rules
|
||||
try:
|
||||
cors = s3_client.get_bucket_cors(Bucket=default_storage.bucket_name)
|
||||
self.stdout.write(self.style.SUCCESS("CORS Configuration:"))
|
||||
self.stdout.write(json.dumps(cors, indent=2))
|
||||
except ClientError as e:
|
||||
if e.response["Error"]["Code"] == "NoSuchCORSConfiguration":
|
||||
self.stdout.write(
|
||||
self.style.WARNING("No CORS configuration found for this bucket.")
|
||||
)
|
||||
else:
|
||||
self.stderr.write(self.style.ERROR(f"Error fetching CORS config: {e}"))
|
||||
@@ -0,0 +1,58 @@
|
||||
"""Custom middleware(s) for the project."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from urllib.parse import unquote
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import MiddlewareNotUsed
|
||||
|
||||
# We are importing posthog here, but it will only be used if the POSTHOG_KEY is set in settings.
|
||||
# The settings are checked before any attempt to use posthog.
|
||||
try:
|
||||
import posthog
|
||||
except ImportError:
|
||||
posthog = None
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PostHogMiddleware:
|
||||
"""
|
||||
This middleware is used to alias the user's distinct_id from the PostHog cookie
|
||||
with their email address when they are authenticated. This allows us to track
|
||||
users across different sessions and devices.
|
||||
"""
|
||||
|
||||
def __init__(self, get_response):
|
||||
"""
|
||||
Initialize the middleware and disable it if PostHog is not configured.
|
||||
"""
|
||||
if posthog is None or not settings.POSTHOG_KEY:
|
||||
raise MiddlewareNotUsed("POSTHOG_KEY must be set in settings to use PostHogMiddleware.")
|
||||
self.get_response = get_response
|
||||
|
||||
def __call__(self, request):
|
||||
"""
|
||||
Process the request to handle the PostHog alias.
|
||||
"""
|
||||
if posthog is not None and settings.POSTHOG_KEY:
|
||||
posthog_cookie = request.COOKIES.get(f"ph_{posthog.project_api_key}_posthog")
|
||||
if posthog_cookie:
|
||||
try:
|
||||
cookie_dict = json.loads(unquote(posthog_cookie))
|
||||
if (
|
||||
cookie_dict.get("distinct_id")
|
||||
and request.user
|
||||
and request.user.is_authenticated
|
||||
):
|
||||
posthog.alias(cookie_dict["distinct_id"], request.user.email)
|
||||
except (json.JSONDecodeError, KeyError):
|
||||
# If the cookie is malformed or doesn't contain the expected
|
||||
# keys, we can't do anything with it, so we ignore it.
|
||||
logger.warning("Malformed PostHog cookie: %s", posthog_cookie)
|
||||
|
||||
response = self.get_response(request)
|
||||
|
||||
return response
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user