Compare commits
60 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0bdee3025b | |||
| b6449addb4 | |||
| f3680b6905 | |||
| 5676ce68c0 | |||
| 50a395c546 | |||
| 69bf2cab5d | |||
| dc61fdce00 | |||
| aa42a9b4d3 | |||
| 5475bcd04e | |||
| b1533c016a | |||
| 9329ee8c90 | |||
| 39bf8f0c2d | |||
| a1ed561204 | |||
| 6dfb9b7328 | |||
| 38ae97aa31 | |||
| 19cf3b2663 | |||
| 83904d8878 | |||
| d3922b7448 | |||
| cdac7cad3b | |||
| 93ee3cd10d | |||
| 91e1c73ccf | |||
| 0493badb12 | |||
| c6283bd8c8 | |||
| e823d21418 | |||
| 22ce90488c | |||
| 8f5419e6ca | |||
| dcec57719f | |||
| 67bb3536d6 | |||
| a020bfa9bf | |||
| 2df761c9a1 | |||
| fe1a065688 | |||
| a5bc974e5d | |||
| 5a3d20f4a9 | |||
| 78b9f11179 | |||
| b33a3e4987 | |||
| c83c8c7da7 | |||
| ee73c7b9cd | |||
| 78a2393383 | |||
| 392eeece3e | |||
| 1f92187dae | |||
| 9426a7e1ae | |||
| e51620a15c | |||
| f1251a3d09 | |||
| 7bc293b8e3 | |||
| bbac17462a | |||
| 7d7ad0bdcd | |||
| eca8fa5ffe | |||
| 5e497b2ccb | |||
| 55400636b6 | |||
| 1901c4d435 | |||
| 095bcaea1a | |||
| fbe9e039cf | |||
| 33a87c3959 | |||
| 18fc3390f3 | |||
| 1d54114a39 | |||
| 7d8b6fc07c | |||
| 2b96ba0597 | |||
| 34cf348f4c | |||
| 0cce897c69 | |||
| 7c8d8e9de7 |
@@ -200,3 +200,22 @@ 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,6 +47,7 @@ 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: .
|
||||
@@ -86,6 +87,7 @@ 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: .
|
||||
|
||||
+72
-7
@@ -8,17 +8,79 @@ and this project adheres to
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Changed
|
||||
|
||||
- 📦️(front) update react
|
||||
|
||||
### Fixed
|
||||
|
||||
- 🐛(e2e) fix test-e2e-chromium
|
||||
- 🐛(back) fix system prompt compatibility with self-hosted models #200
|
||||
- ⚰️(back) remove dead code and unused files
|
||||
|
||||
## [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
|
||||
@@ -30,14 +92,12 @@ and this project adheres to
|
||||
- 🐛(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
|
||||
@@ -47,6 +107,7 @@ and this project adheres to
|
||||
- 📈(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
|
||||
@@ -54,7 +115,6 @@ and this project adheres to
|
||||
- 🐛(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
|
||||
@@ -77,7 +137,7 @@ and this project adheres to
|
||||
- 🎨(front) change list attachment in chat
|
||||
- 🎨(front) move emplacement for attachment
|
||||
- 🎨(ui) retour ui sources files
|
||||
- ✨(ui) fix retour global ui
|
||||
- ✨(ui) fix retour global ui
|
||||
- 🐛(fix) broken staging css
|
||||
- 🎨(alpha) adjustment for alpha version
|
||||
- ✨(ui) delete flex message
|
||||
@@ -112,8 +172,13 @@ and this project adheres to
|
||||
- 💄(chat) add code highlighting for LLM responses #67
|
||||
|
||||
|
||||
[unreleased]: https://github.com/suitenumerique/conversations/compare/v0.0.5...main
|
||||
[0.0.4]: https://github.com/suitenumerique/conversations/releases/v0.0.5
|
||||
[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
|
||||
|
||||
@@ -167,6 +167,7 @@ CMD [\
|
||||
"--host=0.0.0.0",\
|
||||
"--timeout-graceful-shutdown=300",\
|
||||
"--limit-max-requests=20000",\
|
||||
"--lifespan=off",\
|
||||
"conversations.asgi: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 $(cache)
|
||||
@$(COMPOSE_E2E) build frontend openmockllm-mistral $(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
|
||||
@$(COMPOSE_E2E) up --force-recreate -d frontend openmockllm-mistral
|
||||
.PHONY: run-e2e
|
||||
|
||||
status: ## an alias for "docker compose ps"
|
||||
|
||||
@@ -115,6 +115,31 @@ 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:
|
||||
@@ -141,6 +166,18 @@ 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,3 +11,22 @@ 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"
|
||||
|
||||
@@ -7,8 +7,8 @@ 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 <--> Celery --> DB
|
||||
Back --> Cache("Cache (Redis)")
|
||||
Back ----> S3("Minio (S3)")
|
||||
Back -- REST API --> LLM("LLM Providers")
|
||||
```
|
||||
|
||||
@@ -0,0 +1,400 @@
|
||||
# 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
-9
@@ -10,7 +10,6 @@ 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 |
|
||||
@@ -24,12 +23,11 @@ These are the environment variables you can set for the `conversations-backend`
|
||||
| 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 |
|
||||
| ATTACHMENT_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 | |
|
||||
@@ -76,12 +74,14 @@ 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 |
|
||||
| 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 | |
|
||||
| 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) | [] |
|
||||
| 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,7 +9,6 @@ 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
|
||||
|
||||
@@ -7,7 +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 Memcached 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)
|
||||
|
||||
### Test cluster
|
||||
|
||||
@@ -0,0 +1,412 @@
|
||||
# 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
|
||||
|
||||
+20
-20
@@ -14,15 +14,15 @@ 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 (+ 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 |
|
||||
| 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 |
|
||||
|
||||
[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"
|
||||
@@ -58,7 +58,7 @@ Production deployments differ significantly from development environments. The t
|
||||
| Service | Memory | Notes |
|
||||
|----------------------------------|------------|----------------------------------------|
|
||||
| PostgreSQL | **2 GB** | Core database |
|
||||
| Django API (+ Celery) | **1.5 GB** | Backend services |
|
||||
| Django API | **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 +81,16 @@ 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 |
|
||||
| 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) |
|
||||
|
||||
## 6. Sizing Guidelines
|
||||
|
||||
|
||||
+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:
|
||||
|
||||
```javascript
|
||||
```ini
|
||||
FRONTEND_CSS_URL=http://anything/custom-style.css
|
||||
```
|
||||
|
||||
@@ -38,7 +38,7 @@ The footer is configurable from the theme customization file.
|
||||
|
||||
### Settings 🔧
|
||||
|
||||
```shellscript
|
||||
```ini
|
||||
THEME_CUSTOMIZATION_FILE_PATH=<path>
|
||||
```
|
||||
|
||||
@@ -55,10 +55,10 @@ The translations can be partially overridden from the theme customization file.
|
||||
|
||||
### Settings 🔧
|
||||
|
||||
```shellscript
|
||||
```ini
|
||||
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
@@ -0,0 +1,238 @@
|
||||
# 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)
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# 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)
|
||||
|
||||
@@ -0,0 +1,670 @@
|
||||
# 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/)
|
||||
|
||||
@@ -0,0 +1,370 @@
|
||||
# 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,8 +1,10 @@
|
||||
# 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,6 +0,0 @@
|
||||
{
|
||||
"dependencies": {
|
||||
"@ai-sdk/react": "^1.2.12",
|
||||
"@ai-sdk/ui-utils": "^1.2.11"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
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"]
|
||||
@@ -0,0 +1,19 @@
|
||||
[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
|
||||
@@ -3,11 +3,12 @@
|
||||
import json
|
||||
import logging
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
from urllib.parse import urljoin
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
import httpx
|
||||
import requests
|
||||
|
||||
from chat.agent_rag.albert_api_constants import Searches
|
||||
@@ -17,6 +18,169 @@ from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Albert API token limit for document vectorization
|
||||
# We use a conservative chunk size to stay well under the limit
|
||||
ALBERT_MAX_TOKENS = 8192
|
||||
ALBERT_CHUNK_SIZE_TOKENS = 5000 # More conservative chunk size with larger safety margin
|
||||
# Approximate tokens: ~3 characters per token (more conservative estimate for Markdown/Excel)
|
||||
# Markdown and Excel content often have more tokens per character due to formatting
|
||||
ALBERT_CHUNK_SIZE_CHARS = ALBERT_CHUNK_SIZE_TOKENS * 3
|
||||
|
||||
|
||||
def _estimate_tokens(content: str) -> int:
|
||||
"""
|
||||
Estimate the number of tokens in a text string.
|
||||
|
||||
Uses a conservative approximation: ~3 characters per token.
|
||||
This is more conservative than 4 chars/token to account for:
|
||||
- Markdown formatting (headers, lists, tables)
|
||||
- Excel content with special characters
|
||||
- Whitespace and punctuation
|
||||
|
||||
Args:
|
||||
content (str): The text content to estimate.
|
||||
|
||||
Returns:
|
||||
int: Estimated number of tokens.
|
||||
"""
|
||||
return len(content) // 3
|
||||
|
||||
|
||||
def _chunk_content(content: str, max_chars: int = ALBERT_CHUNK_SIZE_CHARS) -> List[str]:
|
||||
"""
|
||||
Split content into chunks that fit within Albert's token limit.
|
||||
|
||||
Attempts to split at paragraph boundaries (double newlines) when possible,
|
||||
otherwise splits at line boundaries, and finally at character boundaries.
|
||||
Validates that each chunk is under the token limit after splitting.
|
||||
|
||||
Args:
|
||||
content (str): The content to chunk.
|
||||
max_chars (int): Maximum characters per chunk (default: ALBERT_CHUNK_SIZE_CHARS).
|
||||
|
||||
Returns:
|
||||
list[str]: List of content chunks, each under the token limit.
|
||||
"""
|
||||
# First check if content fits in one chunk
|
||||
estimated_tokens = _estimate_tokens(content)
|
||||
if estimated_tokens <= ALBERT_CHUNK_SIZE_TOKENS:
|
||||
return [content]
|
||||
|
||||
chunks = []
|
||||
remaining = content
|
||||
|
||||
while len(remaining) > 0:
|
||||
# Check if remaining content fits in one chunk
|
||||
remaining_tokens = _estimate_tokens(remaining)
|
||||
if remaining_tokens <= ALBERT_CHUNK_SIZE_TOKENS:
|
||||
if remaining.strip():
|
||||
chunks.append(remaining.strip())
|
||||
break
|
||||
|
||||
# Need to split - find the best split point
|
||||
# Start with max_chars but may need to reduce if token estimate is too high
|
||||
search_limit = max_chars
|
||||
|
||||
# Try to find a split point that keeps us under token limit
|
||||
# Reduce search limit if needed to ensure token limit is respected
|
||||
while search_limit > 100: # Minimum chunk size
|
||||
# Try to split at paragraph boundary (double newline)
|
||||
split_pos = remaining.rfind("\n\n", 0, search_limit)
|
||||
if split_pos == -1:
|
||||
# Try to split at single newline
|
||||
split_pos = remaining.rfind("\n", 0, search_limit)
|
||||
if split_pos == -1:
|
||||
# Force split at character boundary
|
||||
split_pos = search_limit
|
||||
|
||||
# Validate that this chunk is under token limit
|
||||
chunk_candidate = remaining[:split_pos].strip()
|
||||
if chunk_candidate:
|
||||
chunk_tokens = _estimate_tokens(chunk_candidate)
|
||||
if chunk_tokens <= ALBERT_CHUNK_SIZE_TOKENS:
|
||||
chunks.append(chunk_candidate)
|
||||
remaining = remaining[split_pos:].lstrip()
|
||||
break
|
||||
|
||||
# Chunk too large, reduce search limit and try again
|
||||
search_limit = int(search_limit * 0.8) # Reduce by 20%
|
||||
else:
|
||||
# Fallback: force split at a safe size
|
||||
# This should rarely happen, but ensures we don't get stuck
|
||||
safe_size = min(max_chars, len(remaining))
|
||||
chunk = remaining[:safe_size].strip()
|
||||
if chunk:
|
||||
chunks.append(chunk)
|
||||
remaining = remaining[safe_size:].lstrip()
|
||||
|
||||
# Validate all chunks are under limit and split further if needed
|
||||
validated_chunks = []
|
||||
for chunk_item in chunks:
|
||||
chunk_tokens = _estimate_tokens(chunk_item)
|
||||
if chunk_tokens > ALBERT_MAX_TOKENS:
|
||||
logger.warning(
|
||||
"Chunk still exceeds token limit (%d tokens, max: %d), forcing split further",
|
||||
chunk_tokens,
|
||||
ALBERT_MAX_TOKENS,
|
||||
)
|
||||
# Force split this chunk further using a more conservative size
|
||||
# Use a size that ensures we stay well under the token limit
|
||||
# Target: ~5000 tokens max per chunk (conservative)
|
||||
max_safe_chars = ALBERT_CHUNK_SIZE_TOKENS * 3 # 6000 * 3 = 18000 chars for ~5000 tokens
|
||||
remaining_chunk = chunk_item
|
||||
while len(remaining_chunk) > 0:
|
||||
remaining_tokens = _estimate_tokens(remaining_chunk)
|
||||
if remaining_tokens <= ALBERT_CHUNK_SIZE_TOKENS:
|
||||
if remaining_chunk.strip():
|
||||
validated_chunks.append(remaining_chunk.strip())
|
||||
break
|
||||
|
||||
# Find a safe split point
|
||||
split_pos = min(max_safe_chars, len(remaining_chunk))
|
||||
# Try to split at a line boundary if possible
|
||||
line_split = remaining_chunk.rfind("\n", 0, split_pos)
|
||||
if line_split > max_safe_chars * 0.5: # Only use if it's not too small
|
||||
split_pos = line_split
|
||||
|
||||
sub_chunk = remaining_chunk[:split_pos].strip()
|
||||
if sub_chunk:
|
||||
sub_tokens = _estimate_tokens(sub_chunk)
|
||||
# Double-check this sub-chunk is safe
|
||||
if sub_tokens > ALBERT_MAX_TOKENS:
|
||||
# Still too large, use even smaller size
|
||||
logger.warning(
|
||||
"Sub-chunk still too large (%d tokens), using smaller split",
|
||||
sub_tokens,
|
||||
)
|
||||
split_pos = ALBERT_CHUNK_SIZE_TOKENS * 2 # 12000 chars for ~3000 tokens
|
||||
sub_chunk = remaining_chunk[:split_pos].strip()
|
||||
validated_chunks.append(sub_chunk)
|
||||
remaining_chunk = remaining_chunk[split_pos:].lstrip()
|
||||
else:
|
||||
validated_chunks.append(chunk_item)
|
||||
|
||||
# Final validation - ensure NO chunk exceeds the limit
|
||||
final_chunks = []
|
||||
for chunk in validated_chunks:
|
||||
chunk_tokens = _estimate_tokens(chunk)
|
||||
if chunk_tokens > ALBERT_MAX_TOKENS:
|
||||
logger.error(
|
||||
"CRITICAL: Chunk still exceeds limit after all splitting attempts: %d tokens",
|
||||
chunk_tokens,
|
||||
)
|
||||
# Emergency split: use very conservative size
|
||||
emergency_size = ALBERT_CHUNK_SIZE_TOKENS * 2 # 12000 chars
|
||||
remaining = chunk
|
||||
while len(remaining) > 0:
|
||||
emergency_chunk = remaining[:emergency_size].strip()
|
||||
if emergency_chunk:
|
||||
final_chunks.append(emergency_chunk)
|
||||
remaining = remaining[emergency_size:].lstrip()
|
||||
else:
|
||||
final_chunks.append(chunk)
|
||||
|
||||
return final_chunks
|
||||
|
||||
|
||||
class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attributes
|
||||
"""
|
||||
@@ -32,9 +196,13 @@ 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):
|
||||
def __init__(
|
||||
self,
|
||||
collection_id: Optional[str] = None,
|
||||
read_only_collection_id: Optional[List[str]] = None,
|
||||
):
|
||||
# Initialize any necessary parameters or configurations here
|
||||
super().__init__(collection_id)
|
||||
super().__init__(collection_id, read_only_collection_id)
|
||||
self._base_url = settings.ALBERT_API_URL
|
||||
self._headers = {
|
||||
"Authorization": f"Bearer {settings.ALBERT_API_KEY}",
|
||||
@@ -65,6 +233,27 @@ 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
|
||||
@@ -76,6 +265,18 @@ 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.
|
||||
@@ -132,7 +333,42 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
"""
|
||||
Store the document content in the Albert collection.
|
||||
This method should handle the logic to send the document content to the Albert API.
|
||||
|
||||
If the document is too large (exceeds Albert's token limit), it will be automatically
|
||||
split into multiple chunks and stored as separate documents.
|
||||
|
||||
Args:
|
||||
name (str): The name of the document.
|
||||
content (str): The content of the document in Markdown format.
|
||||
"""
|
||||
# Check if content needs to be chunked
|
||||
estimated_tokens = _estimate_tokens(content)
|
||||
|
||||
if estimated_tokens > ALBERT_MAX_TOKENS:
|
||||
logger.info(
|
||||
"Document '%s' is too large (%d estimated tokens, limit: %d). "
|
||||
"Splitting into chunks.",
|
||||
name,
|
||||
estimated_tokens,
|
||||
ALBERT_MAX_TOKENS,
|
||||
)
|
||||
chunks = _chunk_content(content)
|
||||
logger.info("Split document '%s' into %d chunks", name, len(chunks))
|
||||
|
||||
# Store each chunk as a separate document
|
||||
for i, chunk in enumerate(chunks, start=1):
|
||||
chunk_name = f"{name}_part_{i}" if len(chunks) > 1 else name
|
||||
self._store_single_document(chunk_name, chunk)
|
||||
else:
|
||||
# Document fits within limit, store as-is
|
||||
self._store_single_document(name, content)
|
||||
|
||||
def _store_single_document(self, name: str, content: str) -> None:
|
||||
"""
|
||||
Store a single document chunk in the Albert collection.
|
||||
|
||||
Internal method that performs the actual API call to store one document.
|
||||
|
||||
Args:
|
||||
name (str): The name of the document.
|
||||
content (str): The content of the document in Markdown format.
|
||||
@@ -147,9 +383,88 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
},
|
||||
timeout=settings.ALBERT_API_TIMEOUT,
|
||||
)
|
||||
logger.debug(response.json())
|
||||
logger.debug("Stored document '%s': %s", name, 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.
|
||||
|
||||
If the document is too large (exceeds Albert's token limit), it will be automatically
|
||||
split into multiple chunks and stored as separate documents.
|
||||
|
||||
Args:
|
||||
name (str): The name of the document.
|
||||
content (str): The content of the document in Markdown format.
|
||||
"""
|
||||
# Check if content needs to be chunked
|
||||
estimated_tokens = _estimate_tokens(content)
|
||||
|
||||
if estimated_tokens > ALBERT_MAX_TOKENS:
|
||||
logger.info(
|
||||
"Document '%s' is too large (%d estimated tokens, limit: %d). "
|
||||
"Splitting into chunks.",
|
||||
name,
|
||||
estimated_tokens,
|
||||
ALBERT_MAX_TOKENS,
|
||||
)
|
||||
chunks = _chunk_content(content)
|
||||
logger.info("Split document '%s' into %d chunks", name, len(chunks))
|
||||
|
||||
# Validate chunks before storing
|
||||
for i, chunk in enumerate(chunks, start=1):
|
||||
chunk_tokens = _estimate_tokens(chunk)
|
||||
logger.debug(
|
||||
"Chunk %d/%d: %d chars, ~%d tokens",
|
||||
i,
|
||||
len(chunks),
|
||||
len(chunk),
|
||||
chunk_tokens,
|
||||
)
|
||||
if chunk_tokens > ALBERT_MAX_TOKENS:
|
||||
logger.error(
|
||||
"Chunk %d/%d still exceeds token limit: %d tokens (max: %d)",
|
||||
i,
|
||||
len(chunks),
|
||||
chunk_tokens,
|
||||
ALBERT_MAX_TOKENS,
|
||||
)
|
||||
|
||||
# Store each chunk as a separate document
|
||||
for i, chunk in enumerate(chunks, start=1):
|
||||
chunk_name = f"{name}_part_{i}" if len(chunks) > 1 else name
|
||||
await self._astore_single_document(chunk_name, chunk)
|
||||
else:
|
||||
# Document fits within limit, store as-is
|
||||
await self._astore_single_document(name, content)
|
||||
|
||||
async def _astore_single_document(self, name: str, content: str) -> None:
|
||||
"""
|
||||
Store a single document chunk in the Albert collection.
|
||||
|
||||
Internal method that performs the actual API call to store one document.
|
||||
|
||||
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("Stored document '%s': %s", name, response.json())
|
||||
response.raise_for_status()
|
||||
|
||||
def search(self, query, results_count: int = 4) -> RAGWebResults:
|
||||
"""
|
||||
Perform a search using the Albert API based on the provided query.
|
||||
@@ -161,11 +476,13 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
|
||||
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": [int(self.collection_id)],
|
||||
"collections": collection_ids,
|
||||
"prompt": query,
|
||||
"score_threshold": 0.6,
|
||||
"k": results_count, # Number of chunks to return from the search
|
||||
@@ -190,3 +507,50 @@ 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,9 +1,11 @@
|
||||
"""Implementation of the Albert API for RAG document search."""
|
||||
|
||||
import logging
|
||||
from contextlib import contextmanager
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from io import BytesIO
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
|
||||
from asgiref.sync import sync_to_async
|
||||
|
||||
from chat.agent_rag.constants import RAGWebResults
|
||||
|
||||
@@ -13,11 +15,51 @@ logger = logging.getLogger(__name__)
|
||||
class BaseRagBackend:
|
||||
"""Base class for RAG backends."""
|
||||
|
||||
def __init__(self, collection_id: Optional[str] = None):
|
||||
"""Backend settings."""
|
||||
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.
|
||||
"""
|
||||
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.
|
||||
@@ -25,6 +67,13 @@ 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.
|
||||
@@ -43,8 +92,8 @@ class BaseRagBackend:
|
||||
|
||||
def store_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.
|
||||
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.
|
||||
@@ -52,6 +101,17 @@ 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.
|
||||
@@ -75,12 +135,25 @@ 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:
|
||||
"""
|
||||
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):
|
||||
@@ -92,3 +165,15 @@ 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()
|
||||
|
||||
@@ -190,6 +190,4 @@ class BaseAgent(Agent):
|
||||
|
||||
_tools = [get_pydantic_tools_by_name(tool_name) for tool_name in self.configuration.tools]
|
||||
|
||||
super().__init__(
|
||||
model=_model_instance, system_prompt=_system_prompt, tools=_tools, **kwargs
|
||||
)
|
||||
super().__init__(model=_model_instance, instructions=_system_prompt, tools=_tools, **kwargs)
|
||||
|
||||
@@ -16,7 +16,6 @@ from .base import BaseAgent
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
MOCKED_RESPONSE = """
|
||||
# **Ode to the AI Assistant** 🤖✨
|
||||
|
||||
@@ -102,10 +101,10 @@ class ConversationAgent(BaseAgent):
|
||||
if settings.WARNING_MOCK_CONVERSATION_AGENT:
|
||||
self._model = FunctionModel(stream_function=mocked_agent_model)
|
||||
|
||||
@self.system_prompt
|
||||
@self.instructions
|
||||
def add_the_date() -> str:
|
||||
"""
|
||||
Dynamic system prompt function to add the current date.
|
||||
Dynamic instruction function to add the current date.
|
||||
|
||||
Warning: this will always use the date in the server timezone,
|
||||
not the user's timezone...
|
||||
@@ -113,9 +112,9 @@ class ConversationAgent(BaseAgent):
|
||||
_formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
return f"Today is {_formatted_date}."
|
||||
|
||||
@self.system_prompt
|
||||
@self.instructions
|
||||
def enforce_response_language() -> str:
|
||||
"""Dynamic system prompt function to set the expected language to use."""
|
||||
"""Dynamic instruction function to set the expected language to use."""
|
||||
return f"Answer in {get_language_name(language).lower()}." if language else ""
|
||||
|
||||
def get_web_search_tool_name(self) -> str | None:
|
||||
|
||||
@@ -4,11 +4,6 @@ import dataclasses
|
||||
import logging
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
from asgiref.sync import sync_to_async
|
||||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from .base import BaseAgent
|
||||
|
||||
@@ -26,60 +21,3 @@ class SummarizationAgent(BaseAgent):
|
||||
output_type=str,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@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 hand_off_to_summarization_agent(ctx: RunContext) -> 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.
|
||||
"""
|
||||
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"
|
||||
)
|
||||
text_attachment = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
content_type__startswith="text/",
|
||||
)
|
||||
)
|
||||
|
||||
documents = [await read_document_content(doc) for doc in text_attachment]
|
||||
|
||||
documents_prompt = "\n\n".join(
|
||||
[
|
||||
(f"<document>\n<name>\n{name}\n</name>\n<content>\n{content}\n</content>\n</document>")
|
||||
for name, content 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,
|
||||
metadata={"sources": {doc[0] for doc in documents}},
|
||||
)
|
||||
|
||||
@@ -7,6 +7,7 @@ changes are needed in views.py or tests.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
@@ -25,7 +26,7 @@ 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
|
||||
from pydantic_ai import Agent, InstrumentationSettings, RunContext
|
||||
from pydantic_ai.messages import (
|
||||
BinaryContent,
|
||||
DocumentUrl,
|
||||
@@ -59,7 +60,6 @@ from chat.agents.local_media_url_processors import (
|
||||
update_history_local_urls,
|
||||
update_local_urls,
|
||||
)
|
||||
from chat.agents.summarize import hand_off_to_summarization_agent
|
||||
from chat.ai_sdk_types import (
|
||||
LanguageModelV1Source,
|
||||
SourceUIPart,
|
||||
@@ -72,10 +72,16 @@ from chat.clients.pydantic_ui_message_converter import (
|
||||
ui_message_to_user_content,
|
||||
)
|
||||
from chat.mcp_servers import get_mcp_servers
|
||||
from chat.tools.data_analysis import add_data_analysis_tool
|
||||
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
|
||||
|
||||
# Keep at the top of the file to avoid mocking issues
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
User = get_user_model()
|
||||
@@ -115,7 +121,8 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
self.language = language # might be None
|
||||
self._last_stop_check = 0
|
||||
|
||||
self._store_analytics = settings.LANGFUSE_ENABLED and user.allow_conversation_analytics
|
||||
self._langfuse_available = settings.LANGFUSE_ENABLED
|
||||
self._store_analytics = self._langfuse_available and user.allow_conversation_analytics
|
||||
self.event_encoder = EventEncoder("v4") # Always use v4 for now
|
||||
|
||||
self._support_streaming = True
|
||||
@@ -136,9 +143,16 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
self.conversation_agent = ConversationAgent(
|
||||
model_hrid=self.model_hrid,
|
||||
language=self.language,
|
||||
instrument=self._store_analytics,
|
||||
instrument=InstrumentationSettings(
|
||||
include_binary_content=self._store_analytics,
|
||||
include_content=self._store_analytics,
|
||||
)
|
||||
if self._langfuse_available
|
||||
else False,
|
||||
deps_type=ContextDeps,
|
||||
)
|
||||
add_document_rag_search_tool_from_setting(self.conversation_agent, self.user)
|
||||
add_data_analysis_tool(self.conversation_agent)
|
||||
|
||||
@property
|
||||
def _stop_cache_key(self):
|
||||
@@ -173,7 +187,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._store_analytics:
|
||||
if self._langfuse_available:
|
||||
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))
|
||||
|
||||
@@ -185,7 +199,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
"""Return Vercel-AI-SDK formatted events."""
|
||||
await self._clean()
|
||||
with ExitStack() as stack:
|
||||
if self._store_analytics:
|
||||
if self._langfuse_available:
|
||||
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):
|
||||
@@ -227,6 +241,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
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
|
||||
@@ -240,8 +255,6 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
):
|
||||
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)
|
||||
if not document_store.collection_id:
|
||||
# Create a new collection for the conversation
|
||||
@@ -278,7 +291,24 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
content=document.data,
|
||||
)
|
||||
|
||||
if not document.media_type.startswith("text/"):
|
||||
# Don't convert tabular files (CSV, Excel) to Markdown - keep originals for data_analysis tool
|
||||
# Tabular files are already text-based or can be used directly
|
||||
is_tabular_file = (
|
||||
document.media_type in [
|
||||
"text/csv",
|
||||
"application/csv",
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-excel",
|
||||
"application/excel",
|
||||
]
|
||||
or any(
|
||||
document.identifier.lower().endswith(ext)
|
||||
for ext in [".csv", ".xlsx", ".xls", ".xlsm", ".xlsb"]
|
||||
)
|
||||
)
|
||||
|
||||
# Only convert non-text files that are not tabular files
|
||||
if not document.media_type.startswith("text/") and not is_tabular_file:
|
||||
md_attachment = await models.ChatConversationAttachment.objects.acreate(
|
||||
conversation=self.conversation,
|
||||
uploaded_by=self.user,
|
||||
@@ -352,7 +382,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
return
|
||||
|
||||
# Langfuse settings
|
||||
if self._store_analytics:
|
||||
if self._langfuse_available:
|
||||
langfuse = get_client()
|
||||
langfuse.update_current_trace(
|
||||
session_id=str(self.conversation.pk),
|
||||
@@ -376,8 +406,10 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
self.conversation, input_images, updated_url=image_key_mapping
|
||||
)
|
||||
|
||||
if self._store_analytics:
|
||||
langfuse.update_current_trace(input=user_prompt)
|
||||
if self._langfuse_available:
|
||||
langfuse.update_current_trace(
|
||||
input=user_prompt if self._store_analytics else "REDACTED"
|
||||
)
|
||||
|
||||
usage = {"promptTokens": 0, "completionTokens": 0}
|
||||
|
||||
@@ -409,6 +441,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
await self.parse_input_documents(input_documents)
|
||||
except Exception as exc: # pylint: disable=broad-except
|
||||
@@ -446,7 +479,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
|
||||
if force_web_search:
|
||||
|
||||
@self.conversation_agent.system_prompt
|
||||
@self.conversation_agent.instructions
|
||||
def force_web_search_prompt() -> str:
|
||||
"""Dynamic system prompt function to force web search."""
|
||||
return (
|
||||
@@ -473,6 +506,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
.aexists()
|
||||
)
|
||||
|
||||
|
||||
document_urls = []
|
||||
if not conversation_has_documents and not has_not_pdf_docs:
|
||||
# No documents to process
|
||||
@@ -480,7 +514,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
elif has_not_pdf_docs:
|
||||
add_document_rag_search_tool(self.conversation_agent)
|
||||
|
||||
@self.conversation_agent.system_prompt
|
||||
@self.conversation_agent.instructions
|
||||
def summarization_system_prompt() -> str:
|
||||
return (
|
||||
"When you receive a result from the summarization tool, you MUST return it "
|
||||
@@ -493,13 +527,27 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
"You may add a follow-up question after the summary if needed."
|
||||
)
|
||||
|
||||
@self.conversation_agent.tool
|
||||
async def summarize(ctx) -> ToolReturn:
|
||||
"""
|
||||
Summarize the documents for the user, only when asked for,
|
||||
the documents are in my context.
|
||||
"""
|
||||
return await hand_off_to_summarization_agent(ctx)
|
||||
# Inform the model (system-level) that documents are attached and available
|
||||
@self.conversation_agent.instructions
|
||||
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)
|
||||
|
||||
if not conversation_has_documents and not has_not_pdf_docs:
|
||||
# No documents to process
|
||||
pass
|
||||
elif has_not_pdf_docs:
|
||||
# Already handled above with RAG tool
|
||||
pass
|
||||
else:
|
||||
conversation_documents = [
|
||||
cd
|
||||
@@ -685,7 +733,7 @@ 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._store_analytics
|
||||
if not self._langfuse_available
|
||||
else f"trace-{langfuse.get_current_trace_id()}"
|
||||
)
|
||||
yield events_v4.StartStepPart(
|
||||
@@ -709,9 +757,10 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
|
||||
image_key_mapping=image_key_mapping or None,
|
||||
)
|
||||
|
||||
if self._store_analytics:
|
||||
langfuse.update_current_trace(output=run.result.output)
|
||||
|
||||
if self._langfuse_available:
|
||||
langfuse.update_current_trace(
|
||||
output=run.result.output if self._store_analytics else "REDACTED"
|
||||
)
|
||||
# Vercel finish message
|
||||
yield events_v4.FinishMessagePart(
|
||||
finish_reason=events_v4.FinishReason.STOP,
|
||||
|
||||
@@ -27,9 +27,14 @@ def test_build_pydantic_agent_success_no_tools():
|
||||
"""Test successful agent creation without tools."""
|
||||
agent = ConversationAgent(model_hrid="default-model")
|
||||
assert isinstance(agent, Agent)
|
||||
assert agent._system_prompts == ()
|
||||
|
||||
instructions = agent._instructions
|
||||
assert len(instructions) == 3
|
||||
assert instructions[0] == "You are a helpful assistant"
|
||||
assert instructions[1].__name__ == "add_the_date"
|
||||
assert instructions[2].__name__ == "enforce_response_language"
|
||||
|
||||
assert agent._system_prompts == ("You are a helpful assistant",)
|
||||
assert agent._instructions == []
|
||||
assert isinstance(agent.model, OpenAIChatModel)
|
||||
assert agent.model.model_name == "model-123"
|
||||
assert str(agent.model.client.base_url) == "https://api.llm.com/v1/"
|
||||
@@ -37,6 +42,7 @@ def test_build_pydantic_agent_success_no_tools():
|
||||
assert agent._function_toolset.tools == {}
|
||||
|
||||
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def test_build_pydantic_agent_with_tools(settings):
|
||||
"""Test successful agent creation with tools."""
|
||||
settings.AI_AGENT_TOOLS = ["get_current_weather"]
|
||||
@@ -44,8 +50,14 @@ def test_build_pydantic_agent_with_tools(settings):
|
||||
agent = ConversationAgent(model_hrid="default-model")
|
||||
assert isinstance(agent, Agent)
|
||||
|
||||
assert agent._system_prompts == ("You are a helpful assistant",)
|
||||
assert agent._instructions == []
|
||||
instructions = agent._instructions
|
||||
assert len(instructions) == 3
|
||||
assert instructions[0] == "You are a helpful assistant"
|
||||
assert instructions[1].__name__ == "add_the_date"
|
||||
assert instructions[1]() == "Today is Friday 25/07/2025."
|
||||
assert instructions[2].__name__ == "enforce_response_language"
|
||||
assert instructions[2]() == ""
|
||||
|
||||
assert isinstance(agent.model, OpenAIChatModel)
|
||||
assert agent.model.model_name == "model-123"
|
||||
assert str(agent.model.client.base_url) == "https://api.llm.com/v1/"
|
||||
@@ -56,21 +68,23 @@ def test_build_pydantic_agent_with_tools(settings):
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
def test_add_dynamic_system_prompt():
|
||||
"""
|
||||
Ensure add_the_date and enforce_response_language system prompt are registered
|
||||
Ensure add_the_date and enforce_response_language instructions are registered
|
||||
and returns proper values.
|
||||
"""
|
||||
agent = ConversationAgent(model_hrid="default-model")
|
||||
|
||||
assert len(agent._system_prompt_functions) == 2
|
||||
assert len(agent._system_prompt_functions) == 0
|
||||
|
||||
assert agent._system_prompt_functions[0].function.__name__ == "add_the_date"
|
||||
assert agent._system_prompt_functions[0].function() == "Today is Friday 25/07/2025."
|
||||
|
||||
assert agent._system_prompt_functions[1].function.__name__ == "enforce_response_language"
|
||||
assert agent._system_prompt_functions[1].function() == ""
|
||||
instructions = agent._instructions
|
||||
assert len(instructions) == 3
|
||||
assert instructions[0] == "You are a helpful assistant"
|
||||
assert instructions[1].__name__ == "add_the_date"
|
||||
assert instructions[1]() == "Today is Friday 25/07/2025."
|
||||
assert instructions[2].__name__ == "enforce_response_language"
|
||||
assert instructions[2]() == ""
|
||||
|
||||
agent = ConversationAgent(model_hrid="default-model", language="fr-fr")
|
||||
assert agent._system_prompt_functions[1].function() == "Answer in french."
|
||||
assert agent._instructions[2]() == "Answer in french."
|
||||
|
||||
|
||||
def test_agent_get_web_search_tool_name(settings):
|
||||
|
||||
+66
@@ -0,0 +1,66 @@
|
||||
"""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."
|
||||
)
|
||||
@@ -0,0 +1,270 @@
|
||||
"""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
|
||||
+14
-16
@@ -3,11 +3,11 @@
|
||||
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 (
|
||||
@@ -37,27 +37,22 @@ 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="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
id=result.id, # Use the generated ID
|
||||
role="user",
|
||||
content="Hello!",
|
||||
parts=[TextUIPart(type="text", text="Hello!")],
|
||||
createdAt=timestamp,
|
||||
)
|
||||
result = model_message_to_ui_message(model_message)
|
||||
|
||||
assert result == expected
|
||||
|
||||
|
||||
@@ -65,14 +60,15 @@ 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="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
id=result.id, # Use the generated ID
|
||||
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
|
||||
|
||||
|
||||
@@ -83,8 +79,10 @@ 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="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
id=result.id, # Use the generated ID
|
||||
role="assistant",
|
||||
content="",
|
||||
parts=[
|
||||
@@ -100,7 +98,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
|
||||
|
||||
|
||||
@@ -109,7 +107,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="f0cc3bb5-f207-401b-8281-4cba6202991d", # Mocked UUID
|
||||
id=str(uuid.uuid4()), # not used in comparison
|
||||
role="assistant",
|
||||
content="",
|
||||
parts=[
|
||||
@@ -122,7 +120,7 @@ def test_model_message_to_ui_message_reasoning_full():
|
||||
createdAt=timezone.now(),
|
||||
)
|
||||
result = model_message_to_ui_message(model_message)
|
||||
assert result.id == expected.id
|
||||
assert result.id == IsUUID(4)
|
||||
assert result.role == expected.role
|
||||
assert result.content == expected.content
|
||||
assert result.createdAt == expected.createdAt
|
||||
|
||||
@@ -8,6 +8,7 @@ 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__)
|
||||
@@ -50,6 +51,40 @@ def mock_ai_agent_service_fixture():
|
||||
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"
|
||||
|
||||
@@ -0,0 +1,403 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,472 @@
|
||||
"""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
|
||||
@@ -0,0 +1,154 @@
|
||||
"""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
@@ -0,0 +1,11 @@
|
||||
"""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,8 +1,6 @@
|
||||
"""Common test fixtures for chat conversation endpoint tests."""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
from django.utils import timezone
|
||||
|
||||
@@ -12,14 +10,6 @@ import respx
|
||||
from freezegun import freeze_time
|
||||
|
||||
|
||||
@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():
|
||||
|
||||
@@ -9,6 +9,7 @@ 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
|
||||
|
||||
@@ -23,6 +24,7 @@ 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
|
||||
@@ -88,7 +90,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, mock_uuid4):
|
||||
def test_post_conversation_data_protocol(api_client, mock_openai_stream):
|
||||
"""Test posting messages to a conversation using the 'data' protocol."""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
|
||||
@@ -115,15 +117,29 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream, mock_uu
|
||||
|
||||
# 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
assert mock_openai_stream.called
|
||||
|
||||
# ensure instructions are merged as a system prompt
|
||||
last_request_payload = json.loads(respx.calls.last.request.content)
|
||||
assert last_request_payload["messages"][0] == {
|
||||
"content": (
|
||||
"You are a helpful test assistant :)\n\nToday is Friday 25/07/2025.\n\n"
|
||||
"Answer in english."
|
||||
),
|
||||
"role": "system",
|
||||
}
|
||||
|
||||
chat_conversation.refresh_from_db()
|
||||
assert chat_conversation.ui_messages == [
|
||||
{
|
||||
@@ -137,8 +153,9 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream, mock_uu
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -149,8 +166,9 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream, mock_uu
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -161,35 +179,23 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream, mock_uu
|
||||
parts=[TextUIPart(type="text", text="Hello there")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"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",
|
||||
@@ -210,13 +216,14 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream, mock_uu
|
||||
"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, mock_uuid4):
|
||||
def test_post_conversation_text_protocol(api_client, mock_openai_stream):
|
||||
"""Test posting messages to a conversation using the 'text' protocol."""
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
|
||||
@@ -244,6 +251,15 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream, mock_uu
|
||||
assert response_content == "Hello there"
|
||||
|
||||
assert mock_openai_stream.called
|
||||
# ensure instructions are merged as a system prompt
|
||||
last_request_payload = json.loads(respx.calls.last.request.content)
|
||||
assert last_request_payload["messages"][0] == {
|
||||
"content": (
|
||||
"You are a helpful test assistant :)\n\nToday is Friday 25/07/2025.\n\n"
|
||||
"Answer in english."
|
||||
),
|
||||
"role": "system",
|
||||
}
|
||||
|
||||
chat_conversation.refresh_from_db()
|
||||
assert chat_conversation.ui_messages == [
|
||||
@@ -258,8 +274,9 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream, mock_uu
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -270,8 +287,9 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream, mock_uu
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -282,35 +300,23 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream, mock_uu
|
||||
parts=[TextUIPart(type="text", text="Hello there")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"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",
|
||||
@@ -331,13 +337,14 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream, mock_uu
|
||||
"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, mock_uuid4):
|
||||
def test_post_conversation_with_image(api_client, mock_openai_stream_image):
|
||||
"""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"
|
||||
@@ -375,10 +382,14 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
|
||||
# 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -389,11 +400,12 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
# Check the exact structure expected by the AI service
|
||||
assert body["messages"] == [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"content": (
|
||||
"You are a helpful test assistant :)\n\nToday is Friday 25/07/2025."
|
||||
"\n\nAnswer in english."
|
||||
),
|
||||
"role": "system",
|
||||
},
|
||||
{"content": "Today is Friday 25/07/2025.", "role": "system"},
|
||||
{"content": "Answer in english.", "role": "system"},
|
||||
{
|
||||
"content": [
|
||||
{"text": "Hello, what do you see on this picture?", "type": "text"},
|
||||
@@ -439,8 +451,9 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello, what do you see on this picture?",
|
||||
reasoning=None,
|
||||
@@ -461,8 +474,9 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
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=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I see a cat in the picture.",
|
||||
reasoning=None,
|
||||
@@ -473,29 +487,15 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
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,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"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, what do you see on this picture?",
|
||||
@@ -514,6 +514,7 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -534,13 +535,14 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image, mock
|
||||
"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, mock_uuid4, settings):
|
||||
def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settings):
|
||||
"""Ensure tool calls are correctly forwarded and streamed back."""
|
||||
settings.AI_AGENT_TOOLS = ["get_current_weather"]
|
||||
|
||||
@@ -569,6 +571,10 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
|
||||
# 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'
|
||||
@@ -577,7 +583,7 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
'a:{"toolCallId":"xLDcIljdsDrz0idal7tATWSMm2jhMj47","result":{"location":'
|
||||
'"Paris","temperature":22,"unit":"celsius"}}\n'
|
||||
'0:"The current weather in Paris is nice"\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -587,11 +593,12 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
|
||||
assert body["messages"] == [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"content": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"role": "system",
|
||||
},
|
||||
{"content": "Today is Friday 25/07/2025.", "role": "system"},
|
||||
{"content": "Answer in english.", "role": "system"},
|
||||
{"content": [{"text": "Weather in Paris?", "type": "text"}], "role": "user"},
|
||||
]
|
||||
|
||||
@@ -608,8 +615,9 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -620,8 +628,9 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
parts=[TextUIPart(type="text", text="Weather in Paris?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The current weather in Paris is nice",
|
||||
reasoning=None,
|
||||
@@ -644,35 +653,22 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"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": ["Weather in Paris?"],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "tool_call",
|
||||
@@ -701,9 +697,13 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -715,6 +715,7 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
"tool_name": "get_current_weather",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -737,15 +738,14 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, mock_u
|
||||
"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, mock_uuid4, settings
|
||||
):
|
||||
def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool, settings):
|
||||
"""Ensure tool calls are correctly forwarded and streamed back when failing."""
|
||||
settings.AI_AGENT_TOOLS = []
|
||||
|
||||
@@ -774,6 +774,10 @@ def test_post_conversation_tool_call_fails(
|
||||
|
||||
# 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":'
|
||||
@@ -781,7 +785,7 @@ def test_post_conversation_tool_call_fails(
|
||||
'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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -791,11 +795,12 @@ def test_post_conversation_tool_call_fails(
|
||||
|
||||
assert body["messages"] == [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"content": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in french."
|
||||
),
|
||||
"role": "system",
|
||||
},
|
||||
{"content": "Today is Friday 25/07/2025.", "role": "system"},
|
||||
{"content": "Answer in french.", "role": "system"},
|
||||
{"content": [{"text": "Weather in Paris?", "type": "text"}], "role": "user"},
|
||||
]
|
||||
|
||||
@@ -812,8 +817,9 @@ def test_post_conversation_tool_call_fails(
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -824,8 +830,9 @@ def test_post_conversation_tool_call_fails(
|
||||
parts=[TextUIPart(type="text", text="Weather in Paris?")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I cannot give you an answer to that.",
|
||||
reasoning=None,
|
||||
@@ -848,35 +855,22 @@ def test_post_conversation_tool_call_fails(
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in french."
|
||||
),
|
||||
"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 french.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
{
|
||||
"content": ["Weather in Paris?"],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "tool_call",
|
||||
@@ -905,9 +899,13 @@ def test_post_conversation_tool_call_fails(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in french."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -918,6 +916,7 @@ def test_post_conversation_tool_call_fails(
|
||||
"tool_name": "get_current_weather",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -940,6 +939,7 @@ def test_post_conversation_tool_call_fails(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -972,7 +972,6 @@ 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."""
|
||||
@@ -1017,10 +1016,14 @@ 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1034,7 +1037,6 @@ 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,
|
||||
):
|
||||
@@ -1086,6 +1088,9 @@ 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'
|
||||
@@ -1105,13 +1110,13 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
'0:" sca"\n'
|
||||
'0:"tter"\n'
|
||||
'0:"ing."\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":135}}\n'
|
||||
)
|
||||
|
||||
@@ -1130,8 +1135,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Why the sky is blue?",
|
||||
reasoning=None,
|
||||
@@ -1142,8 +1148,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The sky appears blue due to a phenomenon called Rayleigh scattering.",
|
||||
reasoning=None,
|
||||
@@ -1159,35 +1166,21 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are an amazing assistant.\n\nToday is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are an amazing 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": ["Why the sky is blue?"],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": "2025-07-25T10:36:35.297675Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": "stop",
|
||||
@@ -1215,6 +1208,7 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 135,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -1222,9 +1216,7 @@ def test_post_conversation_data_protocol_no_stream(
|
||||
@freeze_time("2025-07-25T10:36:35.297675Z")
|
||||
@respx.mock
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_conversation_async(
|
||||
api_client, mock_openai_stream, mock_uuid4, monkeypatch, caplog
|
||||
):
|
||||
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")
|
||||
|
||||
@@ -1261,10 +1253,14 @@ async def test_post_conversation_async(
|
||||
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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1283,8 +1279,9 @@ async def test_post_conversation_async(
|
||||
|
||||
assert len(chat_conversation.messages) == 2
|
||||
|
||||
assert chat_conversation.messages[0].id == IsUUID(4)
|
||||
assert chat_conversation.messages[0] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[0].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -1295,8 +1292,9 @@ async def test_post_conversation_async(
|
||||
parts=[TextUIPart(type="text", text="Hello")],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4), # Mocked UUID
|
||||
id=chat_conversation.messages[1].id, # don't test the message ID here
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -1307,35 +1305,22 @@ async def test_post_conversation_async(
|
||||
parts=[TextUIPart(type="text", text="Hello there")],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\nAnswer in english."
|
||||
),
|
||||
"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",
|
||||
@@ -1356,5 +1341,6 @@ async def test_post_conversation_async(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
+127
-119
@@ -14,6 +14,7 @@ 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
|
||||
@@ -32,6 +33,7 @@ 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
|
||||
@@ -149,7 +151,7 @@ def fixture_mock_summarization_agent():
|
||||
super().__init__(**kwargs)
|
||||
self._model = FunctionModel(function=summarization_model) # pylint: disable=protected-access
|
||||
|
||||
with mock.patch("chat.agents.summarize.SummarizationAgent", new=SummarizationAgentMock):
|
||||
with mock.patch("chat.tools.document_summarize.SummarizationAgent", new=SummarizationAgentMock):
|
||||
yield
|
||||
|
||||
|
||||
@@ -214,12 +216,12 @@ def fixture_mock_openai_stream():
|
||||
@responses.activate
|
||||
@respx.mock
|
||||
@freeze_time()
|
||||
def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
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,
|
||||
):
|
||||
"""
|
||||
@@ -273,9 +275,11 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
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'
|
||||
@@ -283,19 +287,22 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
'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":"XXX","url":"sample.pdf","title":null,"providerMetadata":{}}\n'
|
||||
'h:{"sourceType":"url","id":"<mocked_uuid>","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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\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=str_mock_uuid4,
|
||||
id=chat_conversation.messages[0].id,
|
||||
createdAt=timezone.now(),
|
||||
content="What does the document say?",
|
||||
reasoning=None,
|
||||
@@ -305,8 +312,10 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
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=str_mock_uuid4,
|
||||
id=chat_conversation.messages[1].id,
|
||||
createdAt=timezone.now(),
|
||||
content="From the document, I can see that it says 'Hello PDF'.",
|
||||
reasoning=None,
|
||||
@@ -318,7 +327,7 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
ToolInvocationUIPart(
|
||||
type="tool-invocation",
|
||||
toolInvocation=ToolInvocationCall(
|
||||
toolCallId=toolcall_id,
|
||||
toolCallId=chat_conversation.messages[1].parts[0].toolInvocation.toolCallId,
|
||||
toolName="document_search_rag",
|
||||
args={"query": "What does the document say?"},
|
||||
state="call",
|
||||
@@ -330,7 +339,7 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
type="source",
|
||||
source=LanguageModelV1Source(
|
||||
sourceType="url",
|
||||
id=str_mock_uuid4,
|
||||
id=chat_conversation.messages[1].parts[2].source.id,
|
||||
url="sample.pdf",
|
||||
title=None,
|
||||
providerMetadata={},
|
||||
@@ -343,58 +352,34 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
_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": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.\n\n"
|
||||
"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.\n\nWhen 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.\n\n"
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; consider them already "
|
||||
"available via the internal store.",
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"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,
|
||||
},
|
||||
{
|
||||
"content": "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.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"content": ["What does the document say?"],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[1] == {
|
||||
"finish_reason": None,
|
||||
@@ -405,7 +390,7 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
"args": '{"query": "What does the document say?"}',
|
||||
"id": None,
|
||||
"part_kind": "tool-call",
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_call_id": chat_conversation.pydantic_messages[1]["parts"][0]["tool_call_id"],
|
||||
"tool_name": "document_search_rag",
|
||||
}
|
||||
],
|
||||
@@ -423,9 +408,26 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 8,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[2] == {
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.\n\n"
|
||||
"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.\n\nWhen 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.\n\n"
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; consider them already "
|
||||
"available via the internal store."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -439,10 +441,11 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
"metadata": {"sources": ["sample.pdf"]},
|
||||
"part_kind": "tool-return",
|
||||
"timestamp": timezone_now,
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_call_id": chat_conversation.pydantic_messages[2]["parts"][0]["tool_call_id"],
|
||||
"tool_name": "document_search_rag",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[3] == {
|
||||
"finish_reason": None,
|
||||
@@ -469,19 +472,20 @@ def test_post_conversation_with_document_upload( # noqa: PLR0913 # pylint: disa
|
||||
"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( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
def test_post_conversation_with_document_upload_feature_disabled(
|
||||
# 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.
|
||||
@@ -526,13 +530,15 @@ def test_post_conversation_with_document_upload_feature_disabled( # noqa: PLR09
|
||||
|
||||
# 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":150,"completionTokens":25}}\n'
|
||||
)
|
||||
|
||||
# This behavior must be improved in the future to inform the user properly
|
||||
assert "Document upload feature is disabled, ignoring input documents." in caplog.text
|
||||
|
||||
@@ -545,7 +551,6 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
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
|
||||
):
|
||||
@@ -556,6 +561,7 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
pdf_base64 = base64.b64encode(sample_pdf_content.read()).decode("utf-8")
|
||||
|
||||
message = UIMessage(
|
||||
id="1",
|
||||
role="user",
|
||||
@@ -600,29 +606,33 @@ 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":"XXX","url":"sample.pdf.md",'
|
||||
'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'
|
||||
'0:"The document discusses various topics."\n'
|
||||
'f:{"messageId":"XXX"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":201,"completionTokens":13}}\n'
|
||||
).replace("XXX", str_mock_uuid4).replace("pyd_ai_YYY", toolcall_id)
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":287,"completionTokens":19}}\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=str_mock_uuid4,
|
||||
id=chat_conversation.messages[0].id,
|
||||
createdAt=timezone.now(),
|
||||
content="Make a summary of this document.",
|
||||
reasoning=None,
|
||||
@@ -632,8 +642,10 @@ 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=str_mock_uuid4,
|
||||
id=chat_conversation.messages[1].id,
|
||||
createdAt=timezone.now(),
|
||||
content="The document discusses various topics.",
|
||||
reasoning=None,
|
||||
@@ -645,7 +657,7 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
ToolInvocationUIPart(
|
||||
type="tool-invocation",
|
||||
toolInvocation=ToolInvocationCall(
|
||||
toolCallId=toolcall_id,
|
||||
toolCallId=chat_conversation.messages[1].parts[0].toolInvocation.toolCallId,
|
||||
toolName="summarize",
|
||||
args={},
|
||||
state="call",
|
||||
@@ -657,7 +669,7 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
type="source",
|
||||
source=LanguageModelV1Source(
|
||||
sourceType="url",
|
||||
id=str_mock_uuid4,
|
||||
id=chat_conversation.messages[1].parts[2].source.id,
|
||||
url="sample.pdf.md", # might be fixed in the future
|
||||
title=None,
|
||||
providerMetadata={},
|
||||
@@ -670,58 +682,35 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
_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": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.\n\n"
|
||||
"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.\n\nWhen 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.\n\n"
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; consider them already "
|
||||
"available via the internal store."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"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,
|
||||
},
|
||||
{
|
||||
"content": "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.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
{
|
||||
"content": ["Make a summary of this document."],
|
||||
"part_kind": "user-prompt",
|
||||
"timestamp": timezone_now,
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[1] == {
|
||||
"finish_reason": None,
|
||||
@@ -732,7 +721,7 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
"args": "{}",
|
||||
"id": None,
|
||||
"part_kind": "tool-call",
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_call_id": chat_conversation.pydantic_messages[1]["parts"][0]["tool_call_id"],
|
||||
"tool_name": "summarize",
|
||||
}
|
||||
],
|
||||
@@ -750,9 +739,26 @@ 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": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.\n\n"
|
||||
"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.\n\nWhen 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.\n\n"
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; consider them already "
|
||||
"available via the internal store."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -760,10 +766,11 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
"metadata": {"sources": ["sample.pdf.md"]},
|
||||
"part_kind": "tool-return",
|
||||
"timestamp": timezone_now,
|
||||
"tool_call_id": toolcall_id,
|
||||
"tool_call_id": chat_conversation.pydantic_messages[2]["parts"][0]["tool_call_id"],
|
||||
"tool_name": "summarize",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
}
|
||||
assert chat_conversation.pydantic_messages[3] == {
|
||||
"finish_reason": None,
|
||||
@@ -786,4 +793,5 @@ def test_post_conversation_with_document_upload_summarize( # pylint: disable=to
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 6,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
}
|
||||
|
||||
+147
-183
@@ -1,6 +1,8 @@
|
||||
"""Unit tests for chat conversation actions with document URL."""
|
||||
# pylint: disable=too-many-lines
|
||||
|
||||
import uuid
|
||||
|
||||
# pylint: disable=too-many-lines
|
||||
from io import BytesIO
|
||||
|
||||
from django.core.files.storage import default_storage
|
||||
@@ -8,13 +10,13 @@ from django.utils import formats, timezone
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from pydantic_ai import ModelRequest, RequestUsage
|
||||
from pydantic_ai.messages import (
|
||||
DocumentUrl,
|
||||
ModelMessage,
|
||||
ModelResponse,
|
||||
SystemPromptPart,
|
||||
TextPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
@@ -27,6 +29,7 @@ from chat.ai_sdk_types import (
|
||||
UIMessage,
|
||||
)
|
||||
from chat.factories import ChatConversationAttachmentFactory, 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
|
||||
@@ -57,11 +60,11 @@ def fixture_sample_document_content():
|
||||
|
||||
@responses.activate
|
||||
@freeze_time()
|
||||
def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
def test_post_conversation_with_local_pdf_document_url(
|
||||
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
api_client,
|
||||
sample_document_content,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -117,7 +120,7 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
presigned_url = messages[0].parts[3].content[1].url
|
||||
presigned_url = messages[0].parts[0].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
|
||||
@@ -126,11 +129,6 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
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 document?",
|
||||
@@ -142,7 +140,10 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions=f"You are a helpful test assistant :)\n\n{today_promt_date}"
|
||||
"\n\nAnswer in english.",
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "This is a document about a single pixel."
|
||||
@@ -162,20 +163,26 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
|
||||
# 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 == (
|
||||
f'9:{{"toolCallId":"{mock_uuid4}","toolName":"document_parsing",'
|
||||
'9:{"toolCallId":"XXX","toolName":"document_parsing",'
|
||||
'"args":{"documents":[{"identifier":"sample.pdf"}]}}\n'
|
||||
f'a:{{"toolCallId":"{mock_uuid4}","result":{{"state":"done"}}}}\n'
|
||||
'a:{"toolCallId":"XXX","result":{"state":"done"}}\n'
|
||||
'0:"This is a document about a single pixel."\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[0].id,
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this document?",
|
||||
reasoning=None,
|
||||
@@ -189,8 +196,10 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
TextUIPart(type="text", text="What is in this document?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=chat_conversation.messages[1].id,
|
||||
createdAt=timezone.now(),
|
||||
content="This is a document about a single pixel.",
|
||||
reasoning=None,
|
||||
@@ -206,29 +215,14 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
timestamp = timezone.now().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
|
||||
_formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.",
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
"What is in this document?",
|
||||
@@ -245,6 +239,7 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
@@ -271,6 +266,7 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 9,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -278,7 +274,6 @@ def test_post_conversation_with_local_pdf_document_url( # pylint: disable=too-m
|
||||
@freeze_time()
|
||||
def test_post_conversation_with_local_document_wrong_url(
|
||||
api_client,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -287,7 +282,7 @@ def test_post_conversation_with_local_document_wrong_url(
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
document_url = f"/media-key/{mock_uuid4}/sample.pdf"
|
||||
document_url = f"/media-key/{uuid.uuid4()}/sample.pdf"
|
||||
|
||||
message = UIMessage(
|
||||
id="1",
|
||||
@@ -326,10 +321,14 @@ def test_post_conversation_with_local_document_wrong_url(
|
||||
|
||||
# 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 == (
|
||||
f'9:{{"toolCallId":"{mock_uuid4}","toolName":"document_parsing",'
|
||||
'9:{"toolCallId":"XXX","toolName":"document_parsing",'
|
||||
'"args":{"documents":[{"identifier":"sample.pdf"}]}}\n'
|
||||
f'a:{{"toolCallId":"{mock_uuid4}",'
|
||||
'a:{"toolCallId":"XXX",'
|
||||
'"result":{"state":"error","error":"Document '
|
||||
'URL does not belong to the conversation."}}\n'
|
||||
'd:{"finishReason":"error","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
@@ -343,7 +342,6 @@ def test_post_conversation_with_local_document_wrong_url(
|
||||
@freeze_time()
|
||||
def test_post_conversation_with_remote_document_url(
|
||||
api_client,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -391,10 +389,14 @@ def test_post_conversation_with_remote_document_url(
|
||||
|
||||
# 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 == (
|
||||
f'9:{{"toolCallId":"{mock_uuid4}","toolName":"document_parsing",'
|
||||
'9:{"toolCallId":"XXX","toolName":"document_parsing",'
|
||||
'"args":{"documents":[{"identifier":"sample.pdf"}]}}\n'
|
||||
f'a:{{"toolCallId":"{mock_uuid4}",'
|
||||
'a:{"toolCallId":"XXX",'
|
||||
'"result":{"state":"error","error":"External document '
|
||||
'URL are not accepted yet."}}\n'
|
||||
'd:{"finishReason":"error","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
@@ -408,8 +410,6 @@ def test_post_conversation_with_remote_document_url(
|
||||
@freeze_time("2025-10-18T20:48:20.286204Z")
|
||||
def test_post_conversation_with_local_document_url_in_history( # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -417,12 +417,14 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
"""
|
||||
chat_conversation_pk = "0be55da5-8eb7-4dad-aa0f-fea454bd5809"
|
||||
document_url = f"/media-key/{chat_conversation_pk}/sample.pdf"
|
||||
formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
|
||||
chat_conversation = ChatConversationFactory(
|
||||
pk=chat_conversation_pk,
|
||||
owner__language="en-us",
|
||||
messages=[
|
||||
UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=str(uuid.uuid4()),
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this document?",
|
||||
reasoning=None,
|
||||
@@ -437,7 +439,7 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
],
|
||||
),
|
||||
UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=str(uuid.uuid4()),
|
||||
createdAt=timezone.now(),
|
||||
content="This is a document about a single pixel.",
|
||||
reasoning=None,
|
||||
@@ -452,27 +454,11 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
],
|
||||
pydantic_messages=[
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
f"Today is {formatted_date}.\n\n"
|
||||
"Answer in english.",
|
||||
"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 document?",
|
||||
@@ -535,7 +521,7 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
presigned_url = messages[0].parts[3].content[1].url
|
||||
presigned_url = messages[0].parts[0].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
|
||||
@@ -544,18 +530,6 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
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 document?",
|
||||
@@ -567,13 +541,18 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions="You are a helpful test assistant :)\n\n"
|
||||
"Today is Saturday 18/10/2025.\n\n"
|
||||
"Answer in english.",
|
||||
run_id=messages[0].run_id,
|
||||
),
|
||||
ModelResponse(
|
||||
parts=[TextPart(content="This is a document about a single pixel.")],
|
||||
usage=RequestUsage(input_tokens=50, output_tokens=9),
|
||||
model_name="function::agent_model",
|
||||
timestamp=timezone.now(),
|
||||
run_id=messages[1].run_id,
|
||||
),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
@@ -583,7 +562,11 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
)
|
||||
]
|
||||
],
|
||||
instructions="You are a helpful test assistant :)\n\n"
|
||||
"Today is Saturday 18/10/2025.\n\n"
|
||||
"Answer in english.",
|
||||
run_id=messages[2].run_id,
|
||||
),
|
||||
]
|
||||
yield "This is a document of square, very small and nice."
|
||||
@@ -603,17 +586,23 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
|
||||
# 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 a document of square, very small and nice."\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[0].id,
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this document?",
|
||||
reasoning=None,
|
||||
@@ -627,8 +616,10 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
TextUIPart(type="text", text="What is in this document?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=chat_conversation.messages[1].id,
|
||||
createdAt=timezone.now(),
|
||||
content="This is a document about a single pixel.",
|
||||
reasoning=None,
|
||||
@@ -640,8 +631,10 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
TextUIPart(type="text", text="This is a document about a single pixel."),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[2].id == IsUUID(4)
|
||||
assert chat_conversation.messages[2] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=chat_conversation.messages[2].id,
|
||||
createdAt=timezone.now(),
|
||||
content="Give more details about this document.",
|
||||
reasoning=None,
|
||||
@@ -653,8 +646,10 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
TextUIPart(type="text", text="Give more details about this document."),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[3].id == IsUUID(4)
|
||||
assert chat_conversation.messages[3] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=chat_conversation.messages[3].id,
|
||||
createdAt=timezone.now(),
|
||||
content="This is a document of square, very small and nice.",
|
||||
reasoning=None,
|
||||
@@ -667,29 +662,14 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[2]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
"Today is Saturday 18/10/2025.\n\n"
|
||||
"Answer in english.",
|
||||
"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 document?",
|
||||
@@ -706,6 +686,7 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
],
|
||||
# no run_id here
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
@@ -732,9 +713,12 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 9,
|
||||
},
|
||||
# no run_id here
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
"Today is Saturday 18/10/2025.\n\n"
|
||||
"Answer in english.",
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -743,6 +727,7 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
@@ -769,6 +754,7 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 11,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -783,10 +769,10 @@ def test_post_conversation_with_local_document_url_in_history( # pylint: disabl
|
||||
("data.csv", "text/csv"),
|
||||
],
|
||||
)
|
||||
def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
def test_post_conversation_with_local_not_pdf_document_url(
|
||||
# pylint: disable=too-many-arguments,too-many-positional-arguments
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
file_name,
|
||||
content_type,
|
||||
@@ -847,33 +833,6 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
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()),
|
||||
SystemPromptPart(
|
||||
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."
|
||||
),
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
SystemPromptPart(
|
||||
content=(
|
||||
"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."
|
||||
),
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
UserPromptPart(
|
||||
content=[
|
||||
"What is in this document?",
|
||||
@@ -881,7 +840,26 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions=(
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.\n\n"
|
||||
"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.\n\nWhen 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.\n\n"
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; "
|
||||
"consider them already available via the internal store."
|
||||
),
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "This is a document about you."
|
||||
@@ -901,20 +879,26 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
|
||||
# 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 == (
|
||||
f'9:{{"toolCallId":"{mock_uuid4}","toolName":"document_parsing",'
|
||||
'9:{"toolCallId":"XXX","toolName":"document_parsing",'
|
||||
f'"args":{{"documents":[{{"identifier":"{file_name}"}}]}}}}\n'
|
||||
f'a:{{"toolCallId":"{mock_uuid4}","result":{{"state":"done"}}}}\n'
|
||||
'a:{"toolCallId":"XXX","result":{"state":"done"}}\n'
|
||||
'0:"This is a document about you."\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":50,"completionTokens":7}}\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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[0].id,
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this document?",
|
||||
reasoning=None,
|
||||
@@ -926,8 +910,10 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
TextUIPart(type="text", text="What is in this document?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=chat_conversation.messages[1].id,
|
||||
createdAt=timezone.now(),
|
||||
content="This is a document about you.",
|
||||
reasoning=None,
|
||||
@@ -943,53 +929,29 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
timestamp = timezone.now().strftime("%Y-%m-%dT%H:%M:%S.%fZ")
|
||||
_formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": (
|
||||
"You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\n"
|
||||
"Answer in english.\n\n"
|
||||
"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.\n\nWhen 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.\n\n"
|
||||
"[Internal context] User documents are attached to this conversation. "
|
||||
"Do not request re-upload of documents; "
|
||||
"consider them already available via the internal store."
|
||||
),
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
"content": "You are a helpful test assistant :)",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"content": today_promt_date,
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"content": "Answer in english.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"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": timestamp,
|
||||
},
|
||||
{
|
||||
"content": "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.",
|
||||
"dynamic_ref": None,
|
||||
"part_kind": "system-prompt",
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
{
|
||||
"content": [
|
||||
"What is in this document?",
|
||||
@@ -998,6 +960,7 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
"timestamp": timestamp,
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
@@ -1024,5 +987,6 @@ def test_post_conversation_with_local_not_pdf_document_url( # noqa: PLR0913 # p
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 7,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -7,6 +7,7 @@ from django.utils import timezone
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
from dirty_equals import IsUUID
|
||||
from freezegun import freeze_time
|
||||
from rest_framework import status
|
||||
|
||||
@@ -18,6 +19,7 @@ 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
|
||||
@@ -200,7 +202,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, mock_uuid4, history_conversation
|
||||
api_client, mock_openai_stream, 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"
|
||||
@@ -226,10 +228,14 @@ 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -259,8 +265,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -271,8 +278,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -290,7 +298,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, mock_uuid4, history_conversation
|
||||
api_client, mock_openai_stream, 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"
|
||||
@@ -335,8 +343,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello",
|
||||
reasoning=None,
|
||||
@@ -347,8 +356,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -363,7 +373,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, mock_uuid4, history_conversation
|
||||
api_client, mock_openai_stream_image, history_conversation
|
||||
):
|
||||
"""
|
||||
Ensure an image URL is correctly forwarded to the AI service with a conversation with history.
|
||||
@@ -403,10 +413,14 @@ 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -452,8 +466,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello, what do you see on this picture?",
|
||||
reasoning=None,
|
||||
@@ -474,8 +489,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I see a cat in the picture.",
|
||||
reasoning=None,
|
||||
@@ -490,7 +506,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, mock_uuid4, settings, history_conversation
|
||||
api_client, mock_openai_stream_tool, settings, history_conversation
|
||||
):
|
||||
"""
|
||||
Ensure tool calls are correctly forwarded and streamed back with a conversation with history.
|
||||
@@ -521,6 +537,10 @@ 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'
|
||||
@@ -529,7 +549,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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -561,8 +581,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -573,8 +594,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The current weather in Paris is nice",
|
||||
reasoning=None,
|
||||
@@ -606,7 +628,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, mock_uuid4, settings, history_conversation
|
||||
api_client, mock_openai_stream_tool, settings, history_conversation
|
||||
):
|
||||
"""
|
||||
Ensure tool calls are correctly forwarded and streamed back when failing with a
|
||||
@@ -638,6 +660,10 @@ 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'
|
||||
@@ -646,7 +672,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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -678,8 +704,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Weather in Paris?",
|
||||
reasoning=None,
|
||||
@@ -690,8 +717,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I cannot give you an answer to that.",
|
||||
reasoning=None,
|
||||
@@ -891,7 +919,7 @@ def history_conversation_with_tool_fixture():
|
||||
history_timestamp = timezone.now().replace(year=2025, month=6, day=15, hour=10, minute=30)
|
||||
|
||||
# Create a conversation with pre-existing messages including a tool invocation
|
||||
conversation = ChatConversationFactory()
|
||||
conversation = ChatConversationFactory(owner__language="nl-nl")
|
||||
|
||||
# Add previous user and assistant messages with tool invocation
|
||||
conversation.messages = [
|
||||
@@ -1147,7 +1175,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, mock_uuid4, history_conversation_with_image
|
||||
api_client, mock_openai_stream, 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"
|
||||
@@ -1173,10 +1201,14 @@ 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1207,8 +1239,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation_with_image.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="What was in that image again?",
|
||||
reasoning=None,
|
||||
@@ -1219,8 +1252,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation_with_image.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="Hello there",
|
||||
reasoning=None,
|
||||
@@ -1238,7 +1272,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, mock_uuid4, settings, history_conversation_with_tool
|
||||
api_client, mock_openai_stream_tool, 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"]
|
||||
@@ -1266,6 +1300,10 @@ 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'
|
||||
@@ -1274,7 +1312,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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1294,8 +1332,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation_with_tool.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="How about Paris weather?",
|
||||
reasoning=None,
|
||||
@@ -1306,8 +1345,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation_with_tool.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="The current weather in Paris is nice",
|
||||
reasoning=None,
|
||||
@@ -1333,9 +1373,13 @@ 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,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\n"
|
||||
"Answer in dutch.",
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -1344,6 +1388,7 @@ 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] == {
|
||||
@@ -1373,10 +1418,13 @@ 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] == {
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
"Today is Friday 25/07/2025.\n\n"
|
||||
"Answer in dutch.",
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -1388,6 +1436,7 @@ 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] == {
|
||||
@@ -1411,13 +1460,14 @@ 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, mock_uuid4, history_conversation_with_tool
|
||||
api_client, mock_openai_stream_image, 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"
|
||||
@@ -1455,10 +1505,14 @@ 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":0,"completionTokens":0}}\n'
|
||||
)
|
||||
|
||||
@@ -1484,8 +1538,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation_with_tool.messages[4].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="How's the weather in this image?",
|
||||
reasoning=None,
|
||||
@@ -1506,8 +1561,9 @@ 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=str(mock_uuid4), # Mocked UUID
|
||||
id=history_conversation_with_tool.messages[5].id,
|
||||
createdAt=timezone.now(), # Mocked timestamp
|
||||
content="I see a cat in the picture.",
|
||||
reasoning=None,
|
||||
|
||||
+89
-117
@@ -1,15 +1,17 @@
|
||||
"""Unit tests for chat conversation actions with image URL."""
|
||||
|
||||
from django.utils import timezone
|
||||
import uuid
|
||||
|
||||
from django.utils import formats, 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,
|
||||
)
|
||||
@@ -22,6 +24,7 @@ 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
|
||||
@@ -53,7 +56,6 @@ def fixture_sample_image_content():
|
||||
@freeze_time("2025-10-18T20:48:20.286204Z")
|
||||
def test_post_conversation_with_local_image_url(
|
||||
api_client,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -84,22 +86,15 @@ def test_post_conversation_with_local_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/")
|
||||
presigned_url = messages[0].parts[0].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
|
||||
|
||||
formatted_date = formats.date_format(timezone.now(), "l d/m/Y", use_l10n=False)
|
||||
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?",
|
||||
@@ -111,7 +106,10 @@ def test_post_conversation_with_local_image_url(
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions="You are a helpful test assistant :)\n\nToday is "
|
||||
f"{formatted_date}.\n\nAnswer in english.",
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "This is an image of a single pixel."
|
||||
@@ -131,17 +129,23 @@ def test_post_conversation_with_local_image_url(
|
||||
|
||||
# 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'f:{{"messageId":"{mock_uuid4}"}}\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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[0].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
@@ -155,8 +159,10 @@ def test_post_conversation_with_local_image_url(
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
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,
|
||||
@@ -169,29 +175,13 @@ def test_post_conversation_with_local_image_url(
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[0]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\n"
|
||||
"Today is Saturday 18/10/2025.\n\nAnswer in english.",
|
||||
"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?",
|
||||
@@ -208,6 +198,7 @@ def test_post_conversation_with_local_image_url(
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
},
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
@@ -230,6 +221,7 @@ def test_post_conversation_with_local_image_url(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 9,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
@@ -238,7 +230,6 @@ def test_post_conversation_with_local_image_url(
|
||||
def test_post_conversation_with_local_image_wrong_url(
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -247,7 +238,7 @@ def test_post_conversation_with_local_image_wrong_url(
|
||||
chat_conversation = ChatConversationFactory(owner__language="en-us")
|
||||
api_client.force_authenticate(user=chat_conversation.owner)
|
||||
|
||||
image_url = f"/media-key/{mock_uuid4}/sample.png"
|
||||
image_url = f"/media-key/{uuid.uuid4()}/sample.png"
|
||||
|
||||
message = UIMessage(
|
||||
id="1",
|
||||
@@ -272,11 +263,6 @@ def test_post_conversation_with_local_image_wrong_url(
|
||||
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?",
|
||||
@@ -288,7 +274,10 @@ def test_post_conversation_with_local_image_wrong_url(
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions=f"You are a helpful test assistant :)\n\n{today_promt_date}"
|
||||
"\n\nAnswer in english.",
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "cannot read image." # IRL a 400 error would be raised by the LLM
|
||||
@@ -308,9 +297,13 @@ def test_post_conversation_with_local_image_wrong_url(
|
||||
|
||||
# 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'f:{{"messageId":"{mock_uuid4}"}}\n'
|
||||
'f:{"messageId":"<mocked_uuid>"}\n'
|
||||
'd:{"finishReason":"stop","usage":{"promptTokens":50,"completionTokens":4}}\n'
|
||||
)
|
||||
|
||||
@@ -322,7 +315,6 @@ def test_post_conversation_with_local_image_wrong_url(
|
||||
def test_post_conversation_with_remote_image_url(
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -356,11 +348,6 @@ def test_post_conversation_with_remote_image_url(
|
||||
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?",
|
||||
@@ -372,7 +359,10 @@ def test_post_conversation_with_remote_image_url(
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions="You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\nAnswer in english.",
|
||||
run_id=messages[0].run_id,
|
||||
)
|
||||
]
|
||||
yield "This is an image of a single pixel."
|
||||
@@ -392,17 +382,23 @@ def test_post_conversation_with_remote_image_url(
|
||||
|
||||
# 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'f:{{"messageId":"{mock_uuid4}"}}\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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[0].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
@@ -416,8 +412,10 @@ def test_post_conversation_with_remote_image_url(
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
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,
|
||||
@@ -435,7 +433,6 @@ def test_post_conversation_with_remote_image_url(
|
||||
def test_post_conversation_with_local_image_url_in_history(
|
||||
api_client,
|
||||
today_promt_date,
|
||||
mock_uuid4,
|
||||
mock_ai_agent_service,
|
||||
):
|
||||
"""
|
||||
@@ -448,7 +445,7 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
owner__language="en-us",
|
||||
messages=[
|
||||
UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=str(uuid.uuid4()),
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
@@ -463,7 +460,7 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
],
|
||||
),
|
||||
UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
id=str(uuid.uuid4()),
|
||||
createdAt=timezone.now(),
|
||||
content="This is an image of a single pixel.",
|
||||
reasoning=None,
|
||||
@@ -478,27 +475,10 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
],
|
||||
pydantic_messages=[
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": f"You are a helpful test assistant :)\n\n{today_promt_date}"
|
||||
"\n\nAnswer in english.",
|
||||
"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?",
|
||||
@@ -561,7 +541,7 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
)
|
||||
|
||||
async def agent_model(messages: list[ModelMessage], _info: AgentInfo):
|
||||
presigned_url = messages[0].parts[3].content[1].url
|
||||
presigned_url = messages[0].parts[0].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
|
||||
@@ -570,18 +550,6 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
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?",
|
||||
@@ -593,7 +561,9 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
),
|
||||
]
|
||||
],
|
||||
instructions="You are a helpful test assistant :)\n\n"
|
||||
f"{today_promt_date}\n\nAnswer in english.",
|
||||
),
|
||||
ModelResponse(
|
||||
parts=[TextPart(content="This is an image of a single pixel.")],
|
||||
@@ -609,7 +579,10 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
],
|
||||
timestamp=timezone.now(),
|
||||
)
|
||||
]
|
||||
],
|
||||
run_id=messages[2].run_id,
|
||||
instructions="You are a helpful test assistant :)\n\n"
|
||||
"Today is Saturday 18/10/2025.\n\nAnswer in english.",
|
||||
),
|
||||
]
|
||||
yield "This is an image of square, very small and nice."
|
||||
@@ -629,17 +602,23 @@ def test_post_conversation_with_local_image_url_in_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:"This is an image of square, very small and nice."\n'
|
||||
f'f:{{"messageId":"{mock_uuid4}"}}\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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[0].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="What is in this image?",
|
||||
reasoning=None,
|
||||
@@ -653,8 +632,10 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
TextUIPart(type="text", text="What is in this image?"),
|
||||
],
|
||||
)
|
||||
|
||||
assert chat_conversation.messages[1].id == IsUUID(4)
|
||||
assert chat_conversation.messages[1] == UIMessage(
|
||||
id=str(mock_uuid4),
|
||||
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,
|
||||
@@ -666,8 +647,10 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
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=str(mock_uuid4),
|
||||
id=chat_conversation.messages[2].id, # don't test the value directly
|
||||
createdAt=timezone.now(),
|
||||
content="Give more details about this image.",
|
||||
reasoning=None,
|
||||
@@ -679,8 +662,10 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
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=str(mock_uuid4),
|
||||
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,
|
||||
@@ -693,29 +678,13 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
],
|
||||
)
|
||||
|
||||
_run_id = chat_conversation.pydantic_messages[2]["run_id"]
|
||||
assert chat_conversation.pydantic_messages == [
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": f"You are a helpful test assistant :)\n\n{today_promt_date}"
|
||||
"\n\nAnswer in english.",
|
||||
"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?",
|
||||
@@ -756,7 +725,8 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
},
|
||||
},
|
||||
{
|
||||
"instructions": None,
|
||||
"instructions": "You are a helpful test assistant :)\n\nToday is Saturday 18/10/2025."
|
||||
"\n\nAnswer in english.",
|
||||
"kind": "request",
|
||||
"parts": [
|
||||
{
|
||||
@@ -765,6 +735,7 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
"timestamp": "2025-10-18T20:48:20.286204Z",
|
||||
}
|
||||
],
|
||||
"run_id": _run_id,
|
||||
},
|
||||
{
|
||||
"finish_reason": None,
|
||||
@@ -791,5 +762,6 @@ def test_post_conversation_with_local_image_url_in_history(
|
||||
"output_audio_tokens": 0,
|
||||
"output_tokens": 11,
|
||||
},
|
||||
"run_id": _run_id,
|
||||
},
|
||||
]
|
||||
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
"""Test the post_stop_steaming view."""
|
||||
"""Test the post_stop_streaming view."""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -25,6 +25,8 @@ 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)
|
||||
|
||||
@@ -18,18 +18,28 @@ 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=False, prepare=only_if_web_search_enabled
|
||||
web_search_brave,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
),
|
||||
"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
|
||||
web_search_tavily,
|
||||
takes_ctx=False,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
),
|
||||
"web_search_albert_rag": Tool(
|
||||
web_search_albert_rag, takes_ctx=True, prepare=only_if_web_search_enabled
|
||||
web_search_albert_rag,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
max_retries=2,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,873 @@
|
||||
"""Data analysis tool for tabular files (CSV, Excel)."""
|
||||
|
||||
import base64
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from io import BytesIO
|
||||
from typing import Any, Dict
|
||||
|
||||
import matplotlib
|
||||
import numpy as np
|
||||
|
||||
matplotlib.use("Agg") # Use non-interactive backend
|
||||
import matplotlib.pyplot as plt
|
||||
import pandas as pd
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.files.storage import default_storage
|
||||
from django.db.models import Q
|
||||
from asgiref.sync import sync_to_async
|
||||
from pydantic_ai import Agent, RunContext
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.utils import generate_retrieve_policy
|
||||
|
||||
from chat.agents.base import BaseAgent, prepare_custom_model
|
||||
from chat.models import ChatConversationAttachment
|
||||
from chat.tools.exceptions import ModelCannotRetry
|
||||
from chat.tools.utils import last_model_retry_soft_fail
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# MIME types for tabular files
|
||||
TABULAR_MIME_TYPES = [
|
||||
"text/csv",
|
||||
"application/csv",
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-excel",
|
||||
"application/excel",
|
||||
]
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def read_tabular_file(attachment) -> bytes:
|
||||
"""Read tabular file content asynchronously."""
|
||||
with default_storage.open(attachment.key, "rb") as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
def detect_csv_separator(file_data: bytes) -> str:
|
||||
"""
|
||||
Detect the CSV separator by analyzing the first few lines.
|
||||
|
||||
Returns the most likely separator: ',', ';', or '\t'
|
||||
"""
|
||||
# Read first 10KB to analyze
|
||||
sample = file_data[:10240].decode("utf-8", errors="ignore")
|
||||
lines = sample.split("\n")[:10] # First 10 lines
|
||||
|
||||
if not lines:
|
||||
return "," # Default to comma
|
||||
|
||||
# Count occurrences of each separator in the first few lines
|
||||
comma_count = sum(line.count(",") for line in lines)
|
||||
semicolon_count = sum(line.count(";") for line in lines)
|
||||
tab_count = sum(line.count("\t") for line in lines)
|
||||
|
||||
# Return the separator with the highest count
|
||||
if tab_count > comma_count and tab_count > semicolon_count:
|
||||
return "\t"
|
||||
elif semicolon_count > comma_count:
|
||||
return ";"
|
||||
else:
|
||||
return "," # Default to comma
|
||||
|
||||
|
||||
def _convert_to_serializable(obj: Any) -> Any:
|
||||
"""
|
||||
Convert pandas/numpy types to Python native types for JSON serialization.
|
||||
|
||||
Handles:
|
||||
- pandas DataFrame/Series
|
||||
- numpy scalars (int64, float64, etc.)
|
||||
- numpy arrays
|
||||
- pandas Timestamp
|
||||
- Other non-serializable types
|
||||
|
||||
Args:
|
||||
obj: The object to convert.
|
||||
|
||||
Returns:
|
||||
A JSON-serializable version of the object.
|
||||
"""
|
||||
|
||||
# Handle pandas DataFrame
|
||||
if isinstance(obj, pd.DataFrame):
|
||||
# Limit number of rows to avoid huge responses
|
||||
if len(obj) > 1000:
|
||||
obj = obj.head(1000)
|
||||
logger.warning("Result truncated to 1000 rows")
|
||||
return obj.to_dict(orient="records")
|
||||
|
||||
# Handle pandas Series
|
||||
if isinstance(obj, pd.Series):
|
||||
# Convert Series to dict, handling index
|
||||
result_dict = obj.to_dict()
|
||||
# Convert any numpy/pandas types in the values
|
||||
return {str(k): _convert_to_serializable(v) for k, v in result_dict.items()}
|
||||
|
||||
# Handle numpy scalars
|
||||
if isinstance(obj, (np.integer, np.floating)):
|
||||
return obj.item() # Convert to Python native int/float
|
||||
|
||||
# Handle numpy arrays
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
|
||||
# Handle pandas Timestamp
|
||||
if isinstance(obj, pd.Timestamp):
|
||||
return obj.isoformat()
|
||||
|
||||
# Handle lists and tuples - recursively convert elements
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [_convert_to_serializable(item) for item in obj]
|
||||
|
||||
# Handle dicts - recursively convert values
|
||||
if isinstance(obj, dict):
|
||||
return {str(k): _convert_to_serializable(v) for k, v in obj.items()}
|
||||
|
||||
# Handle None, bool, int, float, str - these are already serializable
|
||||
if obj is None or isinstance(obj, (bool, int, float, str)):
|
||||
return obj
|
||||
|
||||
# Fallback: try to convert to string
|
||||
try:
|
||||
return str(obj)
|
||||
except Exception:
|
||||
logger.warning("Could not serialize object of type %s, returning None", type(obj))
|
||||
return None
|
||||
|
||||
|
||||
def _is_valid_excel_file(file_data: bytes, file_name: str) -> bool:
|
||||
"""
|
||||
Check if the file data appears to be a valid Excel file.
|
||||
|
||||
XLSX files are ZIP archives, so they should start with ZIP signature (PK\x03\x04).
|
||||
XLS files have a different signature.
|
||||
"""
|
||||
if not file_data:
|
||||
return False
|
||||
|
||||
file_lower = file_name.lower()
|
||||
|
||||
# Check for XLSX (ZIP-based) signature
|
||||
if file_lower.endswith((".xlsx", ".xlsm", ".xlsb")):
|
||||
# XLSX files are ZIP archives, should start with PK\x03\x04
|
||||
return file_data[:4] == b"PK\x03\x04"
|
||||
|
||||
# Check for XLS (OLE2) signature
|
||||
if file_lower.endswith(".xls"):
|
||||
# XLS files are OLE2 compound documents, should start with specific signature
|
||||
# Common signatures: 0xD0CF11E0 (OLE2) or 0x504B0304 (sometimes saved as ZIP)
|
||||
return (
|
||||
file_data[:4] == b"\xd0\xcf\x11\xe0" # OLE2 signature
|
||||
or file_data[:4] == b"PK\x03\x04" # Sometimes XLS are actually ZIP
|
||||
)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def load_dataframe(file_data: bytes, content_type: str, file_name: str) -> Dict[str, pd.DataFrame]:
|
||||
"""
|
||||
Load tabular file into pandas DataFrames.
|
||||
|
||||
Returns a dictionary mapping sheet/table names to DataFrames.
|
||||
For CSV files, uses 'default' as the key.
|
||||
For Excel files, uses sheet names as keys.
|
||||
"""
|
||||
try:
|
||||
# Handle CSV files - also accept text/plain if file extension is .csv
|
||||
if content_type in ["text/csv", "application/csv"] or (
|
||||
content_type == "text/plain" and file_name.lower().endswith(".csv")
|
||||
):
|
||||
# Detect the separator
|
||||
separator = detect_csv_separator(file_data)
|
||||
logger.debug("Detected CSV separator: %r", separator)
|
||||
|
||||
# Read CSV with detected separator
|
||||
df = pd.read_csv(
|
||||
BytesIO(file_data),
|
||||
sep=separator,
|
||||
on_bad_lines="skip", # Skip problematic lines
|
||||
engine="python", # More flexible parser
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
if df.empty:
|
||||
raise ValueError("CSV file appears to be empty or could not be parsed")
|
||||
|
||||
return {"default": df}
|
||||
elif content_type in [
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.ms-excel",
|
||||
"application/excel",
|
||||
] or file_name.lower().endswith((".xlsx", ".xls", ".xlsm", ".xlsb")):
|
||||
# Validate Excel file format before attempting to read
|
||||
if not _is_valid_excel_file(file_data, file_name):
|
||||
logger.warning(
|
||||
"File '%s' does not appear to be a valid Excel file. "
|
||||
"File size: %d bytes, First bytes: %r",
|
||||
file_name,
|
||||
len(file_data),
|
||||
file_data[:20] if len(file_data) >= 20 else file_data,
|
||||
)
|
||||
raise ValueError(
|
||||
f"File '{file_name}' does not appear to be a valid Excel file. "
|
||||
"It may be corrupted or in an unsupported format."
|
||||
)
|
||||
|
||||
file_lower = file_name.lower()
|
||||
dataframes = {}
|
||||
|
||||
# Try different engines based on file extension
|
||||
if file_lower.endswith(".xls"):
|
||||
# Old Excel format - try xlrd engine
|
||||
try:
|
||||
logger.debug("Attempting to read .xls file with xlrd engine")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="xlrd")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except Exception as xlrd_error:
|
||||
logger.warning("Failed to read .xls with xlrd: %s", xlrd_error)
|
||||
# Fallback: try openpyxl (sometimes .xls files are actually .xlsx)
|
||||
try:
|
||||
logger.debug("Trying openpyxl as fallback for .xls file")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="openpyxl")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except Exception as openpyxl_error:
|
||||
logger.error("Failed to read .xls with both engines: %s", openpyxl_error)
|
||||
raise ValueError(
|
||||
f"Failed to read Excel file '{file_name}': "
|
||||
f"xlrd error: {xlrd_error}, openpyxl error: {openpyxl_error}"
|
||||
) from openpyxl_error
|
||||
else:
|
||||
# XLSX format - try openpyxl first
|
||||
try:
|
||||
logger.debug("Attempting to read Excel file with openpyxl engine")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="openpyxl")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except Exception as openpyxl_error:
|
||||
logger.warning("Failed to read with openpyxl: %s", openpyxl_error)
|
||||
# Try calamine engine if available (faster and more robust)
|
||||
try:
|
||||
logger.debug("Trying calamine engine as fallback")
|
||||
excel_file = pd.ExcelFile(BytesIO(file_data), engine="calamine")
|
||||
dataframes = {
|
||||
sheet_name: excel_file.parse(sheet_name)
|
||||
for sheet_name in excel_file.sheet_names
|
||||
}
|
||||
except ImportError:
|
||||
logger.debug("calamine engine not available")
|
||||
raise ValueError(
|
||||
f"Failed to read Excel file '{file_name}' with openpyxl: {openpyxl_error}. "
|
||||
"The file may be corrupted or in an unsupported format."
|
||||
) from openpyxl_error
|
||||
except Exception as calamine_error:
|
||||
logger.error("Failed to read with calamine: %s", calamine_error)
|
||||
raise ValueError(
|
||||
f"Failed to read Excel file '{file_name}': "
|
||||
f"openpyxl error: {openpyxl_error}, calamine error: {calamine_error}"
|
||||
) from calamine_error
|
||||
|
||||
if not dataframes:
|
||||
raise ValueError(f"Excel file '{file_name}' contains no readable sheets")
|
||||
|
||||
logger.info(
|
||||
"Successfully loaded Excel file '%s' with %d sheet(s): %s",
|
||||
file_name,
|
||||
len(dataframes),
|
||||
list(dataframes.keys()),
|
||||
)
|
||||
return dataframes
|
||||
else:
|
||||
raise ValueError(f"Unsupported content type: {content_type}")
|
||||
except Exception as e:
|
||||
logger.error("Error loading tabular file: %s", e, exc_info=True)
|
||||
raise ModelCannotRetry(f"Failed to load file: {str(e)}") from e
|
||||
|
||||
|
||||
def generate_metadata(dataframes: Dict[str, pd.DataFrame], file_name: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Generate metadata about the tabular file.
|
||||
|
||||
Returns:
|
||||
Dictionary containing:
|
||||
- sheets: List of sheet/table names
|
||||
- schemas: Dictionary mapping sheet names to their schemas
|
||||
- row_counts: Dictionary mapping sheet names to row counts
|
||||
- column_info: Dictionary mapping sheet names to column information
|
||||
"""
|
||||
metadata = {
|
||||
"file_name": file_name,
|
||||
"sheets": list(dataframes.keys()),
|
||||
"schemas": {},
|
||||
"row_counts": {},
|
||||
"column_info": {},
|
||||
}
|
||||
|
||||
for sheet_name, df in dataframes.items():
|
||||
# Schema: column names and types
|
||||
metadata["schemas"][sheet_name] = {
|
||||
col: str(dtype) for col, dtype in df.dtypes.items()
|
||||
}
|
||||
|
||||
# Row count
|
||||
metadata["row_counts"][sheet_name] = len(df)
|
||||
|
||||
# Column info: name, type, sample values, null counts
|
||||
metadata["column_info"][sheet_name] = {}
|
||||
for col in df.columns:
|
||||
col_info = {
|
||||
"type": str(df[col].dtype),
|
||||
"null_count": int(df[col].isna().sum()),
|
||||
"unique_count": int(df[col].nunique()),
|
||||
}
|
||||
# Add sample values (non-null)
|
||||
sample_values = df[col].dropna().head(5).tolist()
|
||||
if sample_values:
|
||||
col_info["sample_values"] = [str(v) for v in sample_values]
|
||||
metadata["column_info"][sheet_name][col] = col_info
|
||||
|
||||
return metadata
|
||||
|
||||
|
||||
async def generate_query(
|
||||
user_query: str, metadata: Dict[str, Any], query_agent: BaseAgent, ctx: RunContext
|
||||
) -> str:
|
||||
"""
|
||||
Use an LLM agent to generate a pandas query from user query and file metadata.
|
||||
"""
|
||||
metadata_str = json.dumps(metadata, indent=2)
|
||||
|
||||
prompt = f"""You are a data analysis assistant. Given a user query and file metadata, generate a Python pandas query to answer the question.
|
||||
|
||||
File metadata:
|
||||
{metadata_str}
|
||||
|
||||
User query: {user_query}
|
||||
|
||||
Generate a Python code snippet that:
|
||||
1. Uses pandas operations (filter, groupby, aggregate, etc.)
|
||||
2. Works with the dataframes loaded in memory (available as 'dataframes' dict)
|
||||
3. Assigns the final result to a variable named 'result'
|
||||
4. Handles the specific sheet/table if multiple sheets exist
|
||||
5. ALWAYS handles NaN/NA values in boolean conditions using .notna() or .fillna() before filtering
|
||||
6. If the user asks for a plot/graph/chart, create it using matplotlib and save to 'plot_image' variable as base64
|
||||
|
||||
IMPORTANT RULES:
|
||||
- The code MUST assign the final result to a variable named 'result'
|
||||
- When filtering with conditions, ALWAYS check for NaN first: df[df['col'].notna() & (df['col'] > value)]
|
||||
- Use .dropna() if you need to remove rows with missing values
|
||||
- Use .fillna() if you need to replace missing values
|
||||
- If plotting: use plt (already imported), create the plot, convert to base64:
|
||||
```python
|
||||
plt.figure(figsize=(10, 6))
|
||||
# ... your plot code ...
|
||||
buf = BytesIO()
|
||||
plt.savefig(buf, format='png')
|
||||
buf.seek(0)
|
||||
plot_image = base64.b64encode(buf.getvalue()).decode('utf-8')
|
||||
plt.close()
|
||||
```
|
||||
NOTE: Do NOT use import statements - plt, base64, BytesIO are already available.
|
||||
|
||||
Return ONLY the Python code, without markdown formatting or explanations. The code should be executable and use variables:
|
||||
- 'dataframes': dict mapping sheet names to DataFrames
|
||||
- Sheet names available: {', '.join(metadata['sheets'])}
|
||||
|
||||
Example format (without plot):
|
||||
df = dataframes['default']
|
||||
df = df[df['column'].notna()] # Remove NaN values first
|
||||
result = df[df['column'] > 100].groupby('category').sum()
|
||||
|
||||
Example format (with plot):
|
||||
df = dataframes['default']
|
||||
plt.figure(figsize=(10, 6))
|
||||
plt.plot(df.index, df['close'])
|
||||
plt.xlabel('Index')
|
||||
plt.ylabel('Close')
|
||||
plt.title('Close vs Index')
|
||||
buf = BytesIO()
|
||||
plt.savefig(buf, format='png')
|
||||
buf.seek(0)
|
||||
plot_image = base64.b64encode(buf.getvalue()).decode('utf-8')
|
||||
plt.close()
|
||||
result = "Plot generated successfully. The plot image has been saved and is available in the tool response."
|
||||
|
||||
IMPORTANT:
|
||||
- Do NOT use import statements in the code. All necessary modules (pd, plt, np, base64, BytesIO) are already available. Do NOT use anything else than these modules.
|
||||
- When returning the result text, mention that a plot was generated and will be available in the response, but do NOT include the URL in the text - the system will handle displaying it.
|
||||
|
||||
Generate the query code:"""
|
||||
|
||||
try:
|
||||
response = await query_agent.run(prompt, usage=ctx.usage)
|
||||
query_code = response.output.strip()
|
||||
|
||||
# Extract code from markdown code blocks if present
|
||||
if "```python" in query_code:
|
||||
query_code = query_code.split("```python")[1].split("```")[0].strip()
|
||||
elif "```" in query_code:
|
||||
query_code = query_code.split("```")[1].split("```")[0].strip()
|
||||
|
||||
return query_code
|
||||
except Exception as e:
|
||||
logger.error("Error generating query: %s", e, exc_info=True)
|
||||
raise ModelRetry("Failed to generate query. Please try rephrasing your question.") from e
|
||||
|
||||
|
||||
@sync_to_async
|
||||
def execute_query(query_code: str, dataframes: Dict[str, pd.DataFrame]) -> Any:
|
||||
"""
|
||||
Execute the generated pandas query safely.
|
||||
|
||||
Note: Uses exec() in a restricted environment. The query code is generated
|
||||
by an LLM based on file metadata, so it should be relatively safe, but
|
||||
we restrict the available builtins and globals.
|
||||
"""
|
||||
try:
|
||||
# Pre-process dataframes to handle common issues
|
||||
processed_dataframes = {}
|
||||
for name, df in dataframes.items():
|
||||
# Make a copy to avoid modifying original
|
||||
df_processed = df.copy()
|
||||
# Replace common NaN representations
|
||||
df_processed = df_processed.replace(["", " ", "nan", "NaN", "None", "null"], pd.NA)
|
||||
processed_dataframes[name] = df_processed
|
||||
|
||||
# Create a safe execution environment
|
||||
safe_globals = {
|
||||
"pd": pd,
|
||||
"plt": plt,
|
||||
"np": np,
|
||||
"base64": base64,
|
||||
"BytesIO": BytesIO,
|
||||
"dataframes": processed_dataframes,
|
||||
"__builtins__": {
|
||||
"len": len,
|
||||
"str": str,
|
||||
"int": int,
|
||||
"float": float,
|
||||
"bool": bool,
|
||||
"list": list,
|
||||
"dict": dict,
|
||||
"set": set,
|
||||
"tuple": tuple,
|
||||
"range": range,
|
||||
"sum": sum,
|
||||
"max": max,
|
||||
"min": min,
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
},
|
||||
}
|
||||
|
||||
# Clean up query code - remove any import statements that might cause issues
|
||||
# Split by lines and filter out import statements
|
||||
lines = query_code.split("\n")
|
||||
cleaned_lines = [
|
||||
line
|
||||
for line in lines
|
||||
if not line.strip().startswith("import ") and not line.strip().startswith("from ")
|
||||
]
|
||||
query_code = "\n".join(cleaned_lines)
|
||||
|
||||
# Execute the query in a restricted namespace
|
||||
local_vars = {}
|
||||
exec(query_code, safe_globals, local_vars) # noqa: S102
|
||||
|
||||
# Get the result - check if 'result' variable exists, otherwise try 'df'
|
||||
if "result" in local_vars:
|
||||
result = local_vars["result"]
|
||||
elif "df" in local_vars:
|
||||
result = local_vars["df"]
|
||||
else:
|
||||
# If no explicit result variable, get the last expression
|
||||
# This is a fallback - ideally the LLM should assign to 'result'
|
||||
raise ValueError("Query must assign result to 'result' variable")
|
||||
|
||||
# Check if a plot was generated
|
||||
plot_image = None
|
||||
if "plot_image" in local_vars:
|
||||
plot_image = local_vars["plot_image"]
|
||||
logger.info("Plot image generated")
|
||||
|
||||
# Convert result to a serializable format
|
||||
result = _convert_to_serializable(result)
|
||||
|
||||
return {"result": result, "plot_image": plot_image}
|
||||
except Exception as e:
|
||||
logger.error("Error executing query: %s", e, exc_info=True)
|
||||
# Provide more helpful error message
|
||||
error_msg = str(e)
|
||||
if "NaN" in error_msg or "NA" in error_msg:
|
||||
error_msg = (
|
||||
f"{error_msg}. "
|
||||
"The query may need to handle missing values (NaN/NA) using .notna() or .dropna() before filtering."
|
||||
)
|
||||
raise ModelCannotRetry(f"Failed to execute query: {error_msg}") from e
|
||||
|
||||
|
||||
@last_model_retry_soft_fail
|
||||
async def data_analysis(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Analyze tabular data files (CSV, Excel) based on user query.
|
||||
Can also generate plots/graphs/charts.
|
||||
|
||||
This tool:
|
||||
1. Loads the tabular file(s) from attachments
|
||||
2. Generates metadata about the file structure
|
||||
3. Uses an LLM to generate a pandas query based on user query
|
||||
4. Executes the query and returns results
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The user's data analysis question.
|
||||
|
||||
Returns:
|
||||
ToolReturn: Contains the analysis results and metadata.
|
||||
"""
|
||||
try:
|
||||
# Find tabular files in attachments
|
||||
# First, get all attachments for debugging
|
||||
all_attachments = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.all()
|
||||
)
|
||||
logger.info(
|
||||
"All attachments in conversation: %s",
|
||||
[
|
||||
{
|
||||
"file_name": a.file_name,
|
||||
"content_type": a.content_type,
|
||||
"upload_state": a.upload_state,
|
||||
"conversion_from": a.conversion_from,
|
||||
}
|
||||
for a in all_attachments
|
||||
],
|
||||
)
|
||||
|
||||
# Find tabular files - exclude converted files (they have conversion_from set)
|
||||
# First try by content_type
|
||||
tabular_attachments_by_type = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
content_type__in=TABULAR_MIME_TYPES,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
)
|
||||
|
||||
# If no files found by content_type, try by file extension as fallback
|
||||
# (some systems detect CSV as text/plain instead of text/csv)
|
||||
if not tabular_attachments_by_type:
|
||||
csv_extensions = [".csv", ".xlsx", ".xls"]
|
||||
all_ready_attachments = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
)
|
||||
tabular_attachments = [
|
||||
att
|
||||
for att in all_ready_attachments
|
||||
if any(att.file_name.lower().endswith(ext) for ext in csv_extensions)
|
||||
# Exclude Markdown files (converted files have .md extension or content_type text/markdown)
|
||||
and not att.file_name.lower().endswith(".md")
|
||||
and att.content_type != "text/markdown"
|
||||
]
|
||||
if tabular_attachments:
|
||||
logger.info(
|
||||
"Found %d tabular file(s) by extension fallback (content_type was not recognized): %s",
|
||||
len(tabular_attachments),
|
||||
[f"{a.file_name} ({a.content_type})" for a in tabular_attachments],
|
||||
)
|
||||
else:
|
||||
tabular_attachments = tabular_attachments_by_type
|
||||
|
||||
# If still no files found, check if there are converted files that might have originals
|
||||
# This handles the case where an Excel file was converted to Markdown for RAG
|
||||
if not tabular_attachments:
|
||||
# Look for converted files with tabular extensions
|
||||
csv_extensions = [".csv", ".xlsx", ".xls"]
|
||||
converted_attachments = await sync_to_async(list)(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.exclude(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
)
|
||||
|
||||
# For each converted file, try to find the original
|
||||
for converted_att in converted_attachments:
|
||||
if any(converted_att.file_name.lower().endswith(ext) for ext in csv_extensions):
|
||||
# Try to find the original file using conversion_from key
|
||||
original_key = converted_att.conversion_from
|
||||
if original_key:
|
||||
original_attachment = await sync_to_async(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
key=original_key,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
).first
|
||||
)()
|
||||
if original_attachment:
|
||||
logger.info(
|
||||
"Found original file '%s' for converted file '%s'",
|
||||
original_attachment.file_name,
|
||||
converted_att.file_name,
|
||||
)
|
||||
tabular_attachments.append(original_attachment)
|
||||
break
|
||||
|
||||
logger.info(
|
||||
"Found %d tabular attachment(s): %s",
|
||||
len(tabular_attachments),
|
||||
[f"{a.file_name} ({a.content_type})" for a in tabular_attachments],
|
||||
)
|
||||
|
||||
if not tabular_attachments:
|
||||
raise ModelCannotRetry(
|
||||
"No tabular files (CSV or Excel) found in the conversation. "
|
||||
"Please upload a CSV or Excel file first. "
|
||||
"Note: If you uploaded an Excel file that was converted to Markdown for RAG, "
|
||||
"the original file must still be available."
|
||||
)
|
||||
|
||||
# Use the first tabular file
|
||||
attachment = tabular_attachments[0]
|
||||
logger.info("Analyzing file: %s (type: %s)", attachment.file_name, attachment.content_type)
|
||||
|
||||
# Load file data
|
||||
file_data = await read_tabular_file(attachment)
|
||||
|
||||
# Validate that this is actually a valid Excel/CSV file (not a converted Markdown file)
|
||||
# Check if it's an Excel file that should have ZIP signature
|
||||
if attachment.file_name.lower().endswith((".xlsx", ".xls", ".xlsm", ".xlsb")):
|
||||
if not _is_valid_excel_file(file_data, attachment.file_name):
|
||||
logger.warning(
|
||||
"File '%s' does not appear to be a valid Excel file. "
|
||||
"It may be a converted Markdown file. Searching for original...",
|
||||
attachment.file_name,
|
||||
)
|
||||
# Try to find the original file
|
||||
# Look for an attachment with the same name but without conversion_from
|
||||
original_attachment = await sync_to_async(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
file_name=attachment.file_name,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
)
|
||||
.filter(
|
||||
Q(conversion_from__isnull=True) | Q(conversion_from="")
|
||||
)
|
||||
.exclude(pk=attachment.pk)
|
||||
.first
|
||||
)()
|
||||
|
||||
if original_attachment:
|
||||
logger.info(
|
||||
"Found original file '%s' (key: %s), using it instead",
|
||||
original_attachment.file_name,
|
||||
original_attachment.key,
|
||||
)
|
||||
attachment = original_attachment
|
||||
file_data = await read_tabular_file(attachment)
|
||||
elif hasattr(attachment, 'conversion_from') and attachment.conversion_from:
|
||||
# Try to find by key if this file has a conversion_from
|
||||
original_attachment = await sync_to_async(
|
||||
ctx.deps.conversation.attachments.filter(
|
||||
key=attachment.conversion_from,
|
||||
upload_state=AttachmentStatus.READY,
|
||||
).first
|
||||
)()
|
||||
if original_attachment:
|
||||
logger.info(
|
||||
"Found original file via conversion_from: '%s'",
|
||||
original_attachment.file_name,
|
||||
)
|
||||
attachment = original_attachment
|
||||
file_data = await read_tabular_file(attachment)
|
||||
else:
|
||||
raise ModelCannotRetry(
|
||||
f"File '{attachment.file_name}' appears to be a converted Markdown file, "
|
||||
"not the original Excel file. The original file is not available. "
|
||||
"Please re-upload the original Excel file."
|
||||
)
|
||||
else:
|
||||
raise ModelCannotRetry(
|
||||
f"File '{attachment.file_name}' does not appear to be a valid Excel file. "
|
||||
"It may be corrupted or in an unsupported format."
|
||||
)
|
||||
|
||||
# Load into pandas DataFrames
|
||||
dataframes = await load_dataframe(file_data, attachment.content_type, attachment.file_name)
|
||||
|
||||
# Generate metadata
|
||||
metadata = generate_metadata(dataframes, attachment.file_name)
|
||||
logger.debug("File metadata: %s", json.dumps(metadata, indent=2))
|
||||
|
||||
# Generate query using LLM
|
||||
# NOTE:
|
||||
# We intentionally create a "bare" Agent instance here instead of using BaseAgent
|
||||
# with tools enabled. Using BaseAgent would attach all configured tools (including
|
||||
# this data_analysis tool itself), which can cause the model to try to call tools
|
||||
# while we're already inside a tool execution, leading to nested tool calls and
|
||||
# failures like "Failed to generate query. Please try rephrasing your question.".
|
||||
#
|
||||
# Here we reuse the same model configuration as BaseAgent but WITHOUT any tools,
|
||||
# so this internal call is purely text-to-text.
|
||||
llm_config = settings.LLM_CONFIGURATIONS[settings.LLM_DEFAULT_MODEL_HRID]
|
||||
if llm_config.is_custom:
|
||||
model_instance = prepare_custom_model(llm_config)
|
||||
else:
|
||||
# Rely on pydantic-ai's built-in model registry / name inference
|
||||
model_instance = llm_config.model_name
|
||||
|
||||
# Use the same keyword as when using BaseAgent, which forwards to Agent.
|
||||
# On the current pydantic_ai version, the correct kwarg is `output_type`,
|
||||
# not `result_type` (passing `result_type` raises a UserError).
|
||||
query_agent = Agent(model=model_instance, output_type=str)
|
||||
query_code = await generate_query(query, metadata, query_agent, ctx)
|
||||
logger.debug("Generated query: %s", query_code)
|
||||
|
||||
# Execute query
|
||||
try:
|
||||
execution_result = await execute_query(query_code, dataframes)
|
||||
result = execution_result.get("result")
|
||||
plot_image_base64 = execution_result.get("plot_image")
|
||||
except Exception as e:
|
||||
logger.error("Query execution failed: %s", e, exc_info=True)
|
||||
raise ModelRetry(
|
||||
f"Failed to execute the generated query: {str(e)}. "
|
||||
"Please try rephrasing your question."
|
||||
) from e
|
||||
|
||||
# Format result for return
|
||||
return_value = {
|
||||
"query": query,
|
||||
"query_code": query_code,
|
||||
"result": result,
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
# Save plot image to storage if generated
|
||||
plot_url = None
|
||||
plot_attachment = None
|
||||
if plot_image_base64:
|
||||
try:
|
||||
# Decode base64 image
|
||||
plot_image_data = base64.b64decode(plot_image_base64)
|
||||
|
||||
# Generate a unique filename for the plot
|
||||
plot_filename = f"plot_{uuid.uuid4().hex[:8]}.png"
|
||||
plot_key = f"{ctx.deps.conversation.pk}/plots/{plot_filename}"
|
||||
|
||||
# Save to storage
|
||||
await sync_to_async(default_storage.save)(
|
||||
plot_key, BytesIO(plot_image_data)
|
||||
)
|
||||
|
||||
# Create a permanent attachment record in the database
|
||||
plot_attachment = await sync_to_async(ChatConversationAttachment.objects.create)(
|
||||
conversation=ctx.deps.conversation,
|
||||
uploaded_by=ctx.deps.user,
|
||||
key=plot_key,
|
||||
file_name=plot_filename,
|
||||
content_type="image/png",
|
||||
upload_state=AttachmentStatus.READY,
|
||||
size=len(plot_image_data),
|
||||
)
|
||||
|
||||
# Generate presigned URL for immediate access (valid for 1 hour)
|
||||
plot_url = await sync_to_async(generate_retrieve_policy)(plot_key)
|
||||
logger.info(
|
||||
"Plot image saved to storage and database: %s (presigned URL: %s)",
|
||||
plot_key,
|
||||
plot_url[:50] + "..."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Failed to save plot image: %s", e, exc_info=True)
|
||||
# Continue without plot URL if save fails
|
||||
|
||||
if plot_url:
|
||||
# Include both local and presigned URLs
|
||||
return_value["plot_url"] = plot_url # Presigned URL for direct access
|
||||
return_value["plot_local_url"] = f"/media-key/{plot_key}" # Local URL for reference
|
||||
# Include attachment ID for reference
|
||||
if plot_attachment:
|
||||
return_value["plot_attachment_id"] = str(plot_attachment.pk)
|
||||
|
||||
return ToolReturn(
|
||||
return_value=return_value,
|
||||
metadata={"file_name": attachment.file_name, "content_type": attachment.content_type},
|
||||
)
|
||||
|
||||
except (ModelCannotRetry, ModelRetry):
|
||||
# Re-raise these as-is
|
||||
raise
|
||||
except Exception as exc:
|
||||
# Unexpected error - stop and inform user
|
||||
logger.exception("Unexpected error in data_analysis: %s", exc)
|
||||
raise ModelCannotRetry(
|
||||
f"An unexpected error occurred during data analysis: {type(exc).__name__}. "
|
||||
"You must explain this to the user and not try to answer based on your knowledge."
|
||||
) from exc
|
||||
|
||||
|
||||
def add_data_analysis_tool(agent: Agent) -> None:
|
||||
"""Add the data analysis tool to an existing agent."""
|
||||
|
||||
@agent.tool(retries=2)
|
||||
@functools.wraps(data_analysis)
|
||||
async def data_analysis_tool(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Analyze tabular data files (CSV, Excel) based on user query.
|
||||
|
||||
This tool loads tabular files, generates metadata about their structure,
|
||||
uses an LLM to generate a pandas query based on the user's question,
|
||||
executes the query, and returns the results.
|
||||
|
||||
Use this tool when the user asks questions about data in CSV or Excel files,
|
||||
such as:
|
||||
- "What is the average sales by region?"
|
||||
- "Show me the top 10 products by revenue"
|
||||
- "How many records are in this file?"
|
||||
- "Filter data where column X is greater than Y"
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The user's data analysis question.
|
||||
"""
|
||||
# Import here to avoid circular import
|
||||
from chat.tools.data_analysis import data_analysis as _data_analysis
|
||||
|
||||
return await _data_analysis(ctx, query)
|
||||
|
||||
@agent.instructions
|
||||
def data_analysis_instructions() -> str:
|
||||
"""Dynamic system prompt function to add data analysis instructions."""
|
||||
return (
|
||||
"When the user asks questions about data in CSV or Excel files, "
|
||||
"use the data_analysis tool to analyze the data and answer their question. "
|
||||
"The tool will handle loading the file, generating queries, and executing them. "
|
||||
"When a plot is generated, the tool returns a 'plot_url' in the result. "
|
||||
"Use this presigned URL directly in markdown image syntax: . "
|
||||
"Do NOT use local URLs like /media-key/... - always use the presigned URL from plot_url. "
|
||||
"Present the results clearly to the user."
|
||||
)
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
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 term to search the internet for.
|
||||
query (str): The query to search the documents for.
|
||||
"""
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
|
||||
@@ -39,12 +39,10 @@ def add_document_rag_search_tool(agent: Agent) -> None:
|
||||
metadata={"sources": {result.url for result in rag_results.data}},
|
||||
)
|
||||
|
||||
@agent.system_prompt
|
||||
@agent.instructions
|
||||
def document_rag_instructions() -> str:
|
||||
"""Dynamic system prompt function to add RAG instructions if any."""
|
||||
return (
|
||||
"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."
|
||||
"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."
|
||||
)
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
"""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
|
||||
@@ -0,0 +1,12 @@
|
||||
"""Exceptions for tool function retries."""
|
||||
|
||||
from pydantic_ai import ModelRetry
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
@@ -0,0 +1,50 @@
|
||||
"""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,24 +1,42 @@
|
||||
"""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 requests
|
||||
import httpx
|
||||
from asgiref.sync import sync_to_async
|
||||
from pydantic_ai import RunContext, RunUsage
|
||||
from pydantic_ai.exceptions import ModelRetry
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
from trafilatura import extract, fetch_url
|
||||
from trafilatura import extract
|
||||
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__)
|
||||
|
||||
|
||||
def llm_summarize(query: str, text: str) -> str:
|
||||
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:
|
||||
"""
|
||||
Summarize the text using the LLM summarization agent.
|
||||
|
||||
@@ -33,7 +51,7 @@ def llm_summarize(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.
|
||||
he result should be up to 30 lines long.
|
||||
The result should be up to 30 lines long.
|
||||
|
||||
<user request>
|
||||
{query}
|
||||
@@ -44,54 +62,87 @@ he result should be up to 30 lines long.
|
||||
</text to summarize>
|
||||
"""
|
||||
|
||||
result = summarization_agent.run_sync(prompt)
|
||||
result = await summarization_agent.run(prompt)
|
||||
return result.output
|
||||
|
||||
|
||||
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_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
|
||||
|
||||
if (document := cache.get(cache_key)) is not None:
|
||||
|
||||
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:
|
||||
return document
|
||||
|
||||
html = fetch_url(url)
|
||||
document = extract(html, include_comments=False, no_fallback=True) or ""
|
||||
cache.set(cache_key, document, settings.BRAVE_CACHE_TTL)
|
||||
try:
|
||||
# Fetch HTML
|
||||
html = await _fetch_url_async(url, timeout=settings.BRAVE_API_TIMEOUT)
|
||||
|
||||
return document
|
||||
# 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
|
||||
|
||||
|
||||
def _extract_and_summarize_snippets(query: str, url: str) -> List[str]:
|
||||
async def _extract_and_summarize_snippets_async(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).
|
||||
"""
|
||||
# Cache by URL to avoid repeated fetch/extract across calls
|
||||
document = _fetch_and_extract(url)
|
||||
if not document:
|
||||
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
|
||||
return []
|
||||
|
||||
if not settings.BRAVE_SUMMARIZATION_ENABLED:
|
||||
return [document]
|
||||
|
||||
async def _fetch_and_store_async(url: str, document_store) -> None:
|
||||
"""Fetch, extract and store text content from the URL in the document store."""
|
||||
|
||||
try:
|
||||
snippet = llm_summarize(query, document)
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
logger.exception("Summarization failed for %s: %s", url, e)
|
||||
snippet = None
|
||||
document = await _fetch_and_extract_async(url)
|
||||
|
||||
return [snippet] if snippet else []
|
||||
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
|
||||
|
||||
|
||||
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]:
|
||||
async def _query_brave_api_async(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 = {
|
||||
@@ -109,14 +160,53 @@ def _query_brave_api(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()
|
||||
|
||||
json_response = response.json()
|
||||
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()
|
||||
|
||||
# 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", [])
|
||||
# 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
|
||||
|
||||
|
||||
def format_tool_return(raw_search_results: List[dict]) -> ToolReturn:
|
||||
@@ -140,92 +230,132 @@ def format_tool_return(raw_search_results: List[dict]) -> ToolReturn:
|
||||
)
|
||||
|
||||
|
||||
def web_search_brave(query: str) -> ToolReturn:
|
||||
@last_model_retry_soft_fail
|
||||
async def web_search_brave(_ctx: RunContext, 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.
|
||||
"""
|
||||
raw_search_results = _query_brave_api(query)
|
||||
try:
|
||||
raw_search_results = await _query_brave_api_async(query)
|
||||
|
||||
reset_caches() # Clear trafilatura caches to avoid memory bloat/leaks
|
||||
await sync_to_async(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:
|
||||
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"]
|
||||
)
|
||||
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)
|
||||
|
||||
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()
|
||||
# Update raw_search_results with extracted snippets
|
||||
for (idx, _), snippets in zip(to_process, results, strict=True):
|
||||
raw_search_results[idx]["extra_snippets"] = snippets
|
||||
|
||||
return format_tool_return(raw_search_results)
|
||||
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
|
||||
|
||||
|
||||
def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn:
|
||||
@last_model_retry_soft_fail
|
||||
async def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
Search the web for up-to-date information using RAG backend
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The query to search for.
|
||||
"""
|
||||
raw_search_results = _query_brave_api(query)
|
||||
logger.info("Starting web search with RAG backend for query: %s", query)
|
||||
try:
|
||||
raw_search_results = await _query_brave_api_async(query)
|
||||
|
||||
reset_caches() # Clear trafilatura caches to avoid memory bloat/leaks
|
||||
# Clear trafilatura caches in thread pool to avoid blocking
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, reset_caches)
|
||||
|
||||
# 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)
|
||||
# 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)
|
||||
for result in raw_search_results
|
||||
]
|
||||
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)
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
rag_results = document_store.search(
|
||||
query,
|
||||
results_count=settings.BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER,
|
||||
)
|
||||
# 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)
|
||||
|
||||
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
|
||||
# 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
|
||||
# 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
|
||||
|
||||
return format_tool_return(raw_search_results)
|
||||
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
|
||||
|
||||
@@ -221,7 +221,7 @@ class ChatViewSet( # pylint: disable=too-many-ancestors, abstract-method
|
||||
url_path="stop-streaming",
|
||||
url_name="stop-streaming",
|
||||
)
|
||||
def post_stop_steaming(self, request, pk): # pylint: disable=unused-argument
|
||||
def post_stop_streaming(self, request, pk): # pylint: disable=unused-argument
|
||||
"""Handle POST requests to stop streaming the chat conversation.
|
||||
|
||||
This action will put a poison pill in the redis cache to stop any ongoing streaming.
|
||||
@@ -425,7 +425,7 @@ class ChatConversationAttachmentViewSet(
|
||||
if settings.POSTHOG_KEY:
|
||||
posthog.capture(
|
||||
"item_uploaded",
|
||||
distinct_id=request.user.email,
|
||||
distinct_id=str(request.user.pk), # same as set by the frontend
|
||||
properties={
|
||||
"id": attachment.pk,
|
||||
"file_name": attachment.file_name,
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Global fixtures for the backend tests."""
|
||||
|
||||
import posthog
|
||||
import pytest
|
||||
from rest_framework.test import APIClient
|
||||
from urllib3.connectionpool import HTTPConnectionPool
|
||||
@@ -41,3 +42,17 @@ 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
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
{
|
||||
"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"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -313,7 +313,6 @@ 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",
|
||||
]
|
||||
@@ -483,7 +482,11 @@ class Base(BraveSettings, Configuration):
|
||||
THUMBNAIL_ALIASES = {}
|
||||
|
||||
# Session
|
||||
SESSION_ENGINE = "django.contrib.sessions.backends.cache"
|
||||
SESSION_ENGINE = values.Value(
|
||||
"django.contrib.sessions.backends.cache",
|
||||
environ_name="SESSION_ENGINE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
SESSION_CACHE_ALIAS = "default"
|
||||
SESSION_COOKIE_AGE = values.PositiveIntegerValue(
|
||||
default=60 * 60 * 12, environ_name="SESSION_COOKIE_AGE", environ_prefix=None
|
||||
@@ -503,6 +506,7 @@ 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
|
||||
@@ -627,9 +631,6 @@ 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",
|
||||
@@ -716,6 +717,11 @@ 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(
|
||||
@@ -785,6 +791,21 @@ 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(
|
||||
|
||||
@@ -1,12 +1,9 @@
|
||||
"""Conversations core API endpoints"""
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import ValidationError
|
||||
|
||||
from rest_framework import exceptions as drf_exceptions
|
||||
from rest_framework import views as drf_views
|
||||
from rest_framework.decorators import api_view
|
||||
from rest_framework.response import Response
|
||||
|
||||
|
||||
def exception_handler(exc, context):
|
||||
@@ -28,14 +25,3 @@ def exception_handler(exc, context):
|
||||
exc = drf_exceptions.ValidationError(detail=detail)
|
||||
|
||||
return drf_views.exception_handler(exc, context)
|
||||
|
||||
|
||||
# pylint: disable=unused-argument
|
||||
@api_view(["GET"])
|
||||
def get_frontend_configuration(request):
|
||||
"""Returns the frontend configuration dict as configured in settings."""
|
||||
frontend_configuration = {
|
||||
"LANGUAGE_CODE": settings.LANGUAGE_CODE,
|
||||
}
|
||||
frontend_configuration.update(settings.FRONTEND_CONFIGURATION)
|
||||
return Response(frontend_configuration)
|
||||
|
||||
@@ -20,23 +20,3 @@ class UserSerializer(serializers.ModelSerializer):
|
||||
"sub",
|
||||
]
|
||||
read_only_fields = ["id", "email", "full_name", "short_name", "sub"]
|
||||
|
||||
|
||||
class UserLightSerializer(UserSerializer):
|
||||
"""Serialize users with limited fields."""
|
||||
|
||||
id = serializers.SerializerMethodField(read_only=True)
|
||||
email = serializers.SerializerMethodField(read_only=True)
|
||||
|
||||
def get_id(self, _user):
|
||||
"""Return always None. Here to have the same fields than in UserSerializer."""
|
||||
return None
|
||||
|
||||
def get_email(self, _user):
|
||||
"""Return always None. Here to have the same fields than in UserSerializer."""
|
||||
return None
|
||||
|
||||
class Meta:
|
||||
model = models.User
|
||||
fields = ["id", "email", "full_name", "short_name"]
|
||||
read_only_fields = ["id", "email", "full_name", "short_name"]
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
"""Custom authentication classes for the Conversations core app"""
|
||||
|
||||
from django.conf import settings
|
||||
|
||||
from rest_framework.authentication import BaseAuthentication
|
||||
from rest_framework.exceptions import AuthenticationFailed
|
||||
|
||||
|
||||
class ServerToServerAuthentication(BaseAuthentication):
|
||||
"""
|
||||
Custom authentication class for server-to-server requests.
|
||||
Validates the presence and correctness of the Authorization header.
|
||||
"""
|
||||
|
||||
AUTH_HEADER = "Authorization"
|
||||
TOKEN_TYPE = "Bearer" # noqa S105
|
||||
|
||||
def authenticate(self, request):
|
||||
"""
|
||||
Authenticate the server-to-server request by validating the Authorization header.
|
||||
|
||||
This method checks if the Authorization header is present in the request, ensures it
|
||||
contains a valid token with the correct format, and verifies the token against the
|
||||
list of allowed server-to-server tokens. If the header is missing, improperly formatted,
|
||||
or contains an invalid token, an AuthenticationFailed exception is raised.
|
||||
|
||||
Returns:
|
||||
None: If authentication is successful
|
||||
(no user is authenticated for server-to-server requests).
|
||||
|
||||
Raises:
|
||||
AuthenticationFailed: If the Authorization header is missing, malformed,
|
||||
or contains an invalid token.
|
||||
"""
|
||||
auth_header = request.headers.get(self.AUTH_HEADER)
|
||||
if not auth_header:
|
||||
raise AuthenticationFailed("Authorization header is missing.")
|
||||
|
||||
# Validate token format and existence
|
||||
auth_parts = auth_header.split(" ")
|
||||
if len(auth_parts) != 2 or auth_parts[0] != self.TOKEN_TYPE:
|
||||
raise AuthenticationFailed("Invalid authorization header.")
|
||||
|
||||
token = auth_parts[1]
|
||||
if token not in settings.SERVER_TO_SERVER_API_TOKENS:
|
||||
raise AuthenticationFailed("Invalid server-to-server token.")
|
||||
|
||||
# Authentication is successful, but no user is authenticated
|
||||
|
||||
def authenticate_header(self, request):
|
||||
"""Return the WWW-Authenticate header value."""
|
||||
return f"{self.TOKEN_TYPE} realm='Create document server to server'"
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
"""Authentication Backends for the Conversations core app."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.exceptions import SuspiciousOperation
|
||||
@@ -15,19 +14,6 @@ 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.
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from enum import StrEnum
|
||||
|
||||
from django.conf import settings
|
||||
from django.utils.text import slugify
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
@@ -43,3 +44,9 @@ 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,7 +9,6 @@ User = get_user_model()
|
||||
|
||||
try:
|
||||
import posthog
|
||||
from posthog.contexts import get_tags
|
||||
except ImportError:
|
||||
posthog = None
|
||||
|
||||
@@ -39,8 +38,7 @@ def is_feature_enabled(
|
||||
if posthog is not None:
|
||||
return posthog.feature_enabled(
|
||||
frontend_feature_name(feature_name),
|
||||
user.email,
|
||||
person_properties={"$host": get_tags().get("$host")},
|
||||
str(user.pk), # same as set by the frontend
|
||||
)
|
||||
|
||||
# No feature flag manager
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
"""A JSONField for DRF to handle serialization/deserialization."""
|
||||
|
||||
import json
|
||||
|
||||
from rest_framework import serializers
|
||||
|
||||
|
||||
class JSONField(serializers.Field):
|
||||
"""
|
||||
A custom field for handling JSON data.
|
||||
"""
|
||||
|
||||
def to_representation(self, value):
|
||||
"""
|
||||
Convert the JSON string to a Python dictionary for serialization.
|
||||
"""
|
||||
return value
|
||||
|
||||
def to_internal_value(self, data):
|
||||
"""
|
||||
Convert the Python dictionary to a JSON string for deserialization.
|
||||
"""
|
||||
if data is None:
|
||||
return None
|
||||
return json.dumps(data)
|
||||
@@ -2,31 +2,9 @@
|
||||
|
||||
import unicodedata
|
||||
|
||||
import django_filters
|
||||
|
||||
|
||||
def remove_accents(value):
|
||||
"""Remove accents from a string (vélo -> velo)."""
|
||||
return "".join(
|
||||
c for c in unicodedata.normalize("NFD", value) if unicodedata.category(c) != "Mn"
|
||||
)
|
||||
|
||||
|
||||
class AccentInsensitiveCharFilter(django_filters.CharFilter):
|
||||
"""
|
||||
A custom CharFilter that filters on the accent-insensitive value searched.
|
||||
"""
|
||||
|
||||
def filter(self, qs, value):
|
||||
"""
|
||||
Apply the filter to the queryset using the unaccented version of the field.
|
||||
|
||||
Args:
|
||||
qs: The queryset to filter.
|
||||
value: The value to search for in the unaccented field.
|
||||
Returns:
|
||||
A filtered queryset.
|
||||
"""
|
||||
if value:
|
||||
value = remove_accents(value)
|
||||
return super().filter(qs, value)
|
||||
|
||||
@@ -1,58 +0,0 @@
|
||||
"""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
|
||||
@@ -0,0 +1,17 @@
|
||||
# Generated by Django 5.2.8 on 2025-12-01 08:40
|
||||
|
||||
from django.db import migrations, models
|
||||
|
||||
|
||||
class Migration(migrations.Migration):
|
||||
dependencies = [
|
||||
("core", "0002_user_allow_conversation_analytics"),
|
||||
]
|
||||
|
||||
operations = [
|
||||
migrations.AlterField(
|
||||
model_name="user",
|
||||
name="short_name",
|
||||
field=models.CharField(blank=True, max_length=50, null=True, verbose_name="short name"),
|
||||
),
|
||||
]
|
||||
@@ -114,7 +114,7 @@ class User(AbstractBaseUser, BaseModel, auth_models.PermissionsMixin):
|
||||
)
|
||||
|
||||
full_name = models.CharField(_("full name"), max_length=100, null=True, blank=True)
|
||||
short_name = models.CharField(_("short name"), max_length=20, null=True, blank=True)
|
||||
short_name = models.CharField(_("short name"), max_length=50, null=True, blank=True)
|
||||
|
||||
email = models.EmailField(_("identity email address"), blank=True, null=True)
|
||||
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>Generate Document</title>
|
||||
</head>
|
||||
<body>
|
||||
<h2>Generate Document</h2>
|
||||
<form method="post" enctype="multipart/form-data">
|
||||
{% csrf_token %}
|
||||
{{ form.as_p }}
|
||||
<button type="submit">Generate PDF</button>
|
||||
</form>
|
||||
</body>
|
||||
</html>
|
||||
@@ -1,58 +0,0 @@
|
||||
"""Custom template tags for the core application of People."""
|
||||
|
||||
import base64
|
||||
|
||||
from django import template
|
||||
from django.contrib.staticfiles import finders
|
||||
|
||||
from PIL import ImageFile as PillowImageFile
|
||||
|
||||
register = template.Library()
|
||||
|
||||
|
||||
def image_to_base64(file_or_path, close=False):
|
||||
"""
|
||||
Return the src string of the base64 encoding of an image represented by its path
|
||||
or file opened or not.
|
||||
|
||||
Inspired by Django's "get_image_dimensions"
|
||||
"""
|
||||
pil_parser = PillowImageFile.Parser()
|
||||
if hasattr(file_or_path, "read"):
|
||||
file = file_or_path
|
||||
if file.closed and hasattr(file, "open"):
|
||||
file_or_path.open()
|
||||
file_pos = file.tell()
|
||||
file.seek(0)
|
||||
else:
|
||||
try:
|
||||
# pylint: disable=consider-using-with
|
||||
file = open(file_or_path, "rb")
|
||||
except OSError:
|
||||
return ""
|
||||
close = True
|
||||
|
||||
try:
|
||||
image_data = file.read()
|
||||
if not image_data:
|
||||
return ""
|
||||
pil_parser.feed(image_data)
|
||||
if pil_parser.image:
|
||||
mime_type = pil_parser.image.get_format_mimetype()
|
||||
encoded_string = base64.b64encode(image_data)
|
||||
return f"data:{mime_type:s};base64, {encoded_string.decode('utf-8'):s}"
|
||||
return ""
|
||||
finally:
|
||||
if close:
|
||||
file.close()
|
||||
else:
|
||||
file.seek(file_pos)
|
||||
|
||||
|
||||
@register.simple_tag
|
||||
def base64_static(path):
|
||||
"""Return a static file into a base64."""
|
||||
full_path = finders.find(path)
|
||||
if full_path:
|
||||
return image_to_base64(full_path, True)
|
||||
return ""
|
||||
@@ -58,7 +58,7 @@ def test_authentication_getter_existing_user_via_email(django_assert_num_queries
|
||||
|
||||
monkeypatch.setattr(OIDCAuthenticationBackend, "get_userinfo", get_userinfo_mocked)
|
||||
|
||||
with django_assert_num_queries(3): # user by sub, user by mail, update sub
|
||||
with django_assert_num_queries(4): # user by sub, user by mail, unicity check, update sub
|
||||
user = klass.get_or_create_user(access_token="test-token", id_token=None, payload=None)
|
||||
|
||||
assert user == db_user
|
||||
@@ -205,7 +205,7 @@ def test_authentication_getter_existing_user_change_fields_sub(
|
||||
monkeypatch.setattr(OIDCAuthenticationBackend, "get_userinfo", get_userinfo_mocked)
|
||||
|
||||
# One and only one additional update query when a field has changed
|
||||
with django_assert_num_queries(2):
|
||||
with django_assert_num_queries(3): # user by sub, unicity check, update sub
|
||||
authenticated_user = klass.get_or_create_user(
|
||||
access_token="test-token", id_token=None, payload=None
|
||||
)
|
||||
@@ -245,7 +245,7 @@ def test_authentication_getter_existing_user_change_fields_email(
|
||||
monkeypatch.setattr(OIDCAuthenticationBackend, "get_userinfo", get_userinfo_mocked)
|
||||
|
||||
# One and only one additional update query when a field has changed
|
||||
with django_assert_num_queries(3):
|
||||
with django_assert_num_queries(4): # user by sub, user by mail, unicity check, update sub
|
||||
authenticated_user = klass.get_or_create_user(
|
||||
access_token="test-token", id_token=None, payload=None
|
||||
)
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
"""Tests for feature flag helpers."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from unittest.mock import patch
|
||||
|
||||
import posthog
|
||||
import pytest
|
||||
import responses
|
||||
|
||||
from core.factories import UserFactory
|
||||
from core.feature_flags.flags import FeatureToggle
|
||||
@@ -42,19 +45,29 @@ def test_is_feature_enabled_always_disabled(feature_flags):
|
||||
assert is_feature_enabled(user, "document_upload") is False
|
||||
|
||||
|
||||
@patch("core.feature_flags.helpers.posthog")
|
||||
def test_is_feature_enabled_dynamic_posthog_true(mock_posthog, feature_flags):
|
||||
@responses.activate
|
||||
def test_is_feature_enabled_dynamic_posthog_true(feature_flags, settings):
|
||||
"""Test that a dynamic feature returns the value from PostHog when PostHog is available."""
|
||||
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"]
|
||||
|
||||
responses.post(
|
||||
f"{posthog.host}/flags/?v=2", json={"flags": {"web-search": {"enabled": True}}}, status=200
|
||||
)
|
||||
|
||||
feature_flags.web_search = FeatureToggle.DYNAMIC
|
||||
user = UserFactory()
|
||||
|
||||
mock_posthog.feature_enabled.return_value = True
|
||||
assert is_feature_enabled(user, "web_search") is True
|
||||
mock_posthog.feature_enabled.assert_called_once_with(
|
||||
"web-search",
|
||||
user.email,
|
||||
person_properties={"$host": None},
|
||||
)
|
||||
|
||||
request_body = json.loads(responses.calls[0].request.body)
|
||||
assert request_body["distinct_id"] == str(user.pk)
|
||||
assert request_body["flag_keys_to_evaluate"] == ["web-search"]
|
||||
|
||||
posthog.api_key = None
|
||||
posthog.host = None
|
||||
|
||||
|
||||
@patch("core.feature_flags.helpers.posthog")
|
||||
|
||||
@@ -4,9 +4,10 @@ Test config API endpoints in the Conversations core app.
|
||||
|
||||
import json
|
||||
|
||||
from django.test import override_settings
|
||||
from django.test import AsyncClient, override_settings
|
||||
|
||||
import pytest
|
||||
from asgiref.sync import sync_to_async
|
||||
from rest_framework.status import (
|
||||
HTTP_200_OK,
|
||||
)
|
||||
@@ -156,3 +157,52 @@ def test_api_config_with_original_theme_customization(is_authenticated, settings
|
||||
theme_customization = json.load(f)
|
||||
|
||||
assert content["theme_customization"] == theme_customization
|
||||
|
||||
|
||||
@override_settings(
|
||||
CRISP_WEBSITE_ID="123",
|
||||
FRONTEND_CSS_URL="http://testcss/",
|
||||
FRONTEND_THEME="test-theme",
|
||||
MEDIA_BASE_URL="http://testserver/",
|
||||
POSTHOG_KEY={"id": "132456", "host": "https://eu.i.posthog-test.com"},
|
||||
SENTRY_DSN="https://sentry.test/123",
|
||||
THEME_CUSTOMIZATION_FILE_PATH="",
|
||||
RAG_FILES_ACCEPTED_FORMATS=[
|
||||
"application/pdf",
|
||||
"text/plain",
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_authenticated", [False, True])
|
||||
async def test_api_config_async(is_authenticated):
|
||||
"""Anonymous users should be allowed to get the configuration (async client)."""
|
||||
client = AsyncClient()
|
||||
|
||||
if is_authenticated:
|
||||
user = await sync_to_async(factories.UserFactory)()
|
||||
await client.aforce_login(user)
|
||||
|
||||
response = await client.get("/api/v1.0/config/")
|
||||
assert response.status_code == HTTP_200_OK
|
||||
assert response.json() == {
|
||||
"ACTIVATION_REQUIRED": False,
|
||||
"CRISP_WEBSITE_ID": "123",
|
||||
"ENVIRONMENT": "test",
|
||||
"FEATURE_FLAGS": {"document-upload": "enabled", "web-search": "enabled"},
|
||||
"FRONTEND_CSS_URL": "http://testcss/",
|
||||
"FRONTEND_HOMEPAGE_FEATURE_ENABLED": True,
|
||||
"FRONTEND_THEME": "test-theme",
|
||||
"LANGUAGES": [
|
||||
["en-us", "English"],
|
||||
["fr-fr", "Français"],
|
||||
# ["de-de", "Deutsch"],
|
||||
["nl-nl", "Nederlands"],
|
||||
# ["es-es", "Español"],
|
||||
],
|
||||
"LANGUAGE_CODE": "en-us",
|
||||
"MEDIA_BASE_URL": "http://testserver/",
|
||||
"POSTHOG_KEY": {"id": "132456", "host": "https://eu.i.posthog-test.com"},
|
||||
"SENTRY_DSN": "https://sentry.test/123",
|
||||
"theme_customization": {},
|
||||
"chat_upload_accept": "application/pdf,text/plain",
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ msgstr ""
|
||||
"Project-Id-Version: la-suite-conversations\n"
|
||||
"Report-Msgid-Bugs-To: \n"
|
||||
"POT-Creation-Date: 2025-10-20 21:48+0000\n"
|
||||
"PO-Revision-Date: 2025-10-27 08:29\n"
|
||||
"PO-Revision-Date: 2025-12-15 13:49\n"
|
||||
"Last-Translator: \n"
|
||||
"Language-Team: German\n"
|
||||
"Language: de_DE\n"
|
||||
|
||||
@@ -3,7 +3,7 @@ msgstr ""
|
||||
"Project-Id-Version: la-suite-conversations\n"
|
||||
"Report-Msgid-Bugs-To: \n"
|
||||
"POT-Creation-Date: 2025-10-20 21:48+0000\n"
|
||||
"PO-Revision-Date: 2025-10-27 08:29\n"
|
||||
"PO-Revision-Date: 2025-12-15 13:49\n"
|
||||
"Last-Translator: \n"
|
||||
"Language-Team: English\n"
|
||||
"Language: en_US\n"
|
||||
|
||||
@@ -3,7 +3,7 @@ msgstr ""
|
||||
"Project-Id-Version: la-suite-conversations\n"
|
||||
"Report-Msgid-Bugs-To: \n"
|
||||
"POT-Creation-Date: 2025-10-20 21:48+0000\n"
|
||||
"PO-Revision-Date: 2025-10-27 08:29\n"
|
||||
"PO-Revision-Date: 2025-12-15 13:49\n"
|
||||
"Last-Translator: \n"
|
||||
"Language-Team: French\n"
|
||||
"Language: fr_FR\n"
|
||||
@@ -86,7 +86,7 @@ msgstr "A utilisé le code d'activation"
|
||||
|
||||
#: activation_codes/admin.py:293 build/lib/activation_codes/admin.py:293
|
||||
msgid "Add selected users to Brevo waiting list"
|
||||
msgstr "Ajouter les utilisateurs sélectionnés à la liste d'attente de Brevo"
|
||||
msgstr "Ajouter les utilisateurs sélectionnés à la liste d'attente Brevo"
|
||||
|
||||
#: activation_codes/admin.py:314 build/lib/activation_codes/admin.py:314
|
||||
#, python-format
|
||||
@@ -272,7 +272,7 @@ msgstr "Nous n'avons pas pu trouver un utilisateur avec ce sous-groupe mais l'e-
|
||||
|
||||
#: build/lib/core/models.py:99 core/models.py:99
|
||||
msgid "Enter a valid sub. This value may contain only letters, numbers, and @/./+/-/_/: characters."
|
||||
msgstr "Saisissez un sous-groupe valide. Cette valeur ne peut contenir que des lettres, des chiffres et les caractères @/./+/-/_/: uniquement."
|
||||
msgstr "Saisissez un 'sub' valide. Cette valeur ne peut contenir que des lettres, des chiffres et les caractères @/./+/-/_/: uniquement."
|
||||
|
||||
#: build/lib/core/models.py:105 core/models.py:105
|
||||
msgid "sub"
|
||||
|
||||
@@ -3,7 +3,7 @@ msgstr ""
|
||||
"Project-Id-Version: la-suite-conversations\n"
|
||||
"Report-Msgid-Bugs-To: \n"
|
||||
"POT-Creation-Date: 2025-10-20 21:48+0000\n"
|
||||
"PO-Revision-Date: 2025-10-27 08:29\n"
|
||||
"PO-Revision-Date: 2025-12-15 13:49\n"
|
||||
"Last-Translator: \n"
|
||||
"Language-Team: Dutch\n"
|
||||
"Language: nl_NL\n"
|
||||
|
||||
@@ -3,7 +3,7 @@ msgstr ""
|
||||
"Project-Id-Version: la-suite-conversations\n"
|
||||
"Report-Msgid-Bugs-To: \n"
|
||||
"POT-Creation-Date: 2025-10-20 21:48+0000\n"
|
||||
"PO-Revision-Date: 2025-10-27 08:29\n"
|
||||
"PO-Revision-Date: 2025-12-15 13:49\n"
|
||||
"Last-Translator: \n"
|
||||
"Language-Team: Russian\n"
|
||||
"Language: ru_RU\n"
|
||||
|
||||
@@ -3,7 +3,7 @@ msgstr ""
|
||||
"Project-Id-Version: la-suite-conversations\n"
|
||||
"Report-Msgid-Bugs-To: \n"
|
||||
"POT-Creation-Date: 2025-10-20 21:48+0000\n"
|
||||
"PO-Revision-Date: 2025-10-27 08:29\n"
|
||||
"PO-Revision-Date: 2025-12-15 13:49\n"
|
||||
"Last-Translator: \n"
|
||||
"Language-Team: Ukrainian\n"
|
||||
"Language: uk_UA\n"
|
||||
|
||||
+28
-22
@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "conversations"
|
||||
version = "0.0.5"
|
||||
version = "0.0.10"
|
||||
authors = [{ "name" = "DINUM", "email" = "dev@mail.numerique.gouv.fr" }]
|
||||
classifiers = [
|
||||
"Development Status :: 5 - Production/Stable",
|
||||
@@ -27,43 +27,48 @@ requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
"deprecated",
|
||||
"beautifulsoup4==4.14.2",
|
||||
"boto3==1.40.51",
|
||||
"Brotli==1.1.0",
|
||||
"boto3==1.40.73",
|
||||
"Brotli==1.2.0",
|
||||
"django-configurations==2.5.1",
|
||||
"django-cors-headers==4.9.0",
|
||||
"django-countries==7.6.1",
|
||||
"django-countries==8.1.0",
|
||||
"django-filter==25.2",
|
||||
"django-lasuite[all]==0.0.14",
|
||||
"django-lasuite[all]==0.0.18",
|
||||
"django-parler==2.3",
|
||||
"django-pydantic-field==0.3.13",
|
||||
"django-pydantic-field==0.4.0",
|
||||
"django-redis==6.0.0",
|
||||
"django-storages[s3]==1.14.6",
|
||||
"django-timezone-field>=5.1",
|
||||
"django==5.2.7",
|
||||
"django==5.2.9",
|
||||
"djangorestframework==3.16.1",
|
||||
"drf_spectacular==0.28.0",
|
||||
"drf_spectacular==0.29.0",
|
||||
"dockerflow==2024.4.2",
|
||||
"easy_thumbnails==2.10.1",
|
||||
"factory_boy==3.3.3",
|
||||
"gunicorn==23.0.0",
|
||||
"jsonschema==4.25.1",
|
||||
"langfuse==3.6.2",
|
||||
"langfuse==3.10.0",
|
||||
"lxml==5.4.0",
|
||||
"markdown==3.9",
|
||||
"markdown==3.10",
|
||||
"markitdown==0.0.2",
|
||||
"mozilla-django-oidc==4.0.1",
|
||||
"nested-multipart-parser==1.6.0",
|
||||
"posthog==6.7.7",
|
||||
"pydantic==2.12.1",
|
||||
"pydantic-ai-slim[openai,mistral,mcp,evals,logfire]==1.0.18",
|
||||
"psycopg[binary]==3.2.10",
|
||||
"matplotlib==3.9.2",
|
||||
"numpy==2.1.3",
|
||||
"openpyxl==3.1.5",
|
||||
"pandas==2.2.3",
|
||||
"posthog==7.0.0",
|
||||
"pydantic==2.12.4",
|
||||
"pydantic-ai-slim[openai,mistral,mcp,evals,logfire]==1.17.0",
|
||||
"psycopg[binary]==3.2.12",
|
||||
"PyJWT==2.10.1",
|
||||
"python-magic==0.4.27",
|
||||
"redis<6.0.0",
|
||||
"requests==2.32.5",
|
||||
"sentry-sdk==2.41.0",
|
||||
"semchunk==3.2.5",
|
||||
"sentry-sdk==2.44.0",
|
||||
"trafilatura==2.0.0",
|
||||
"uvicorn==0.28.0",
|
||||
"uvicorn==0.38.0",
|
||||
"whitenoise==6.11.0",
|
||||
]
|
||||
|
||||
@@ -75,25 +80,26 @@ dependencies = [
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"dirty-equals==0.10.0",
|
||||
"django-extensions==4.1",
|
||||
"django-test-migrations==1.5.0",
|
||||
"drf-spectacular-sidecar==2025.10.1",
|
||||
"freezegun==1.5.5",
|
||||
"ipdb==0.13.13",
|
||||
"ipython==9.6.0",
|
||||
"pyfakefs==5.10.0",
|
||||
"ipython==9.7.0",
|
||||
"pyfakefs==5.10.2",
|
||||
"pylint-django==2.6.1",
|
||||
"pylint==3.3.9",
|
||||
"pylint-pydantic==0.4.0",
|
||||
"pytest-asyncio==1.2.0",
|
||||
"pylint-pydantic==0.4.1",
|
||||
"pytest-asyncio==1.3.0",
|
||||
"pytest-cov==7.0.0",
|
||||
"pytest-django==4.11.1",
|
||||
"pytest==8.4.2",
|
||||
"pytest==9.0.1",
|
||||
"pytest-icdiff==0.9",
|
||||
"pytest-xdist==3.8.0",
|
||||
"responses==0.25.8",
|
||||
"respx==0.22.0",
|
||||
"ruff==0.14.0",
|
||||
"ruff==0.14.5",
|
||||
"types-requests==2.32.4.20250913",
|
||||
]
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "app-conversations",
|
||||
"version": "0.0.5",
|
||||
"version": "0.0.10",
|
||||
"private": true,
|
||||
"scripts": {
|
||||
"dev": "next dev",
|
||||
@@ -36,16 +36,16 @@
|
||||
"i18next-browser-languagedetector": "8.1.0",
|
||||
"idb": "8.0.3",
|
||||
"lodash": "4.17.21",
|
||||
"lottie-react": "^2.4.1",
|
||||
"luxon": "3.6.1",
|
||||
"micromark-extension-llm-math": "3.1.1-20250610",
|
||||
"next": "15.3.3",
|
||||
"next": "15.3.8",
|
||||
"posthog-js": "1.249.3",
|
||||
"react": "*",
|
||||
"react": "19.2.1",
|
||||
"react-aria-components": "1.9.0",
|
||||
"react-dom": "18.3.1",
|
||||
"react-dom": "19.2.1",
|
||||
"react-i18next": "15.5.2",
|
||||
"react-intersection-observer": "9.16.0",
|
||||
"react-lottie": "^1.2.10",
|
||||
"react-markdown": "10.1.0",
|
||||
"react-select": "5.10.1",
|
||||
"rehype-katex": "7.0.1",
|
||||
|
||||
@@ -3,6 +3,7 @@ import { css } from 'styled-components';
|
||||
|
||||
import { Box, Text } from '@/components';
|
||||
import { Icon } from '@/components/Icon';
|
||||
import { useResponsiveStore } from '@/stores';
|
||||
|
||||
export type ToastType = 'success' | 'error' | 'info' | 'warning';
|
||||
|
||||
@@ -13,6 +14,8 @@ export interface ToastProps {
|
||||
icon?: string;
|
||||
duration?: number;
|
||||
onClose: (id: string) => void;
|
||||
actionLabel?: string;
|
||||
actionHref?: string;
|
||||
}
|
||||
|
||||
const getToastConfig = (type: ToastType) => {
|
||||
@@ -62,11 +65,14 @@ export const Toast = ({
|
||||
icon,
|
||||
duration = 4000,
|
||||
onClose,
|
||||
actionLabel,
|
||||
actionHref,
|
||||
}: ToastProps) => {
|
||||
const [isVisible, setIsVisible] = useState(false);
|
||||
const [isLeaving, setIsLeaving] = useState(false);
|
||||
const config = getToastConfig(type);
|
||||
const iconToUse = icon || config.icon;
|
||||
const { isMobile } = useResponsiveStore();
|
||||
|
||||
useEffect(() => {
|
||||
setIsVisible(true);
|
||||
@@ -102,7 +108,12 @@ export const Toast = ({
|
||||
overflow: hidden;
|
||||
`}
|
||||
>
|
||||
<Box $direction="row" $align="center" $gap="12px">
|
||||
<Box
|
||||
$direction="row"
|
||||
$align="center"
|
||||
$gap="12px"
|
||||
$justify="space-between"
|
||||
>
|
||||
<Icon
|
||||
iconName={iconToUse}
|
||||
$variation="600"
|
||||
@@ -111,16 +122,41 @@ export const Toast = ({
|
||||
color: ${config.color} !important;
|
||||
`}
|
||||
/>
|
||||
<Text
|
||||
$weight="500"
|
||||
$size="14px"
|
||||
$css={css`
|
||||
color: ${config.color} !important;
|
||||
padding: 4px;
|
||||
`}
|
||||
<Box
|
||||
$direction="row"
|
||||
$align="center"
|
||||
$gap="12px"
|
||||
$flex={1}
|
||||
$justify="space-between"
|
||||
>
|
||||
{message}
|
||||
</Text>
|
||||
<Text
|
||||
$weight="500"
|
||||
$size="14px"
|
||||
$css={css`
|
||||
color: ${config.color} !important;
|
||||
padding: 4px;
|
||||
`}
|
||||
>
|
||||
{message}
|
||||
</Text>
|
||||
|
||||
{actionLabel && actionHref && !isMobile && (
|
||||
<a
|
||||
href={actionHref}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
style={{
|
||||
color: config.color,
|
||||
fontWeight: '500',
|
||||
fontSize: '14px',
|
||||
textDecoration: 'underline',
|
||||
whiteSpace: 'nowrap',
|
||||
}}
|
||||
>
|
||||
{actionLabel}
|
||||
</a>
|
||||
)}
|
||||
</Box>
|
||||
</Box>
|
||||
</Box>
|
||||
);
|
||||
|
||||
@@ -28,10 +28,16 @@ const StyledButton = styled(Button)<StyledButtonProps>`
|
||||
border: none;
|
||||
background: none;
|
||||
outline: none;
|
||||
transition: all 0.2s ease-in-out;
|
||||
transition: background 0.2s ease-in-out;
|
||||
font-weight: 500;
|
||||
font-size: 0.938rem;
|
||||
padding: 0;
|
||||
|
||||
&:focus-visible {
|
||||
outline: 2px solid #3e5de7;
|
||||
outline-offset: 2px;
|
||||
}
|
||||
|
||||
${({ $css }) => $css};
|
||||
`;
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { PropsWithChildren, useRef, useState } from 'react';
|
||||
import { PropsWithChildren, useEffect, useRef, useState } from 'react';
|
||||
import { css } from 'styled-components';
|
||||
|
||||
import { Box, BoxButton, BoxProps, DropButton, Icon, Text } from '@/components';
|
||||
@@ -35,8 +35,16 @@ export const DropdownMenu = ({
|
||||
}: PropsWithChildren<DropdownMenuProps>) => {
|
||||
const { spacingsTokens, colorsTokens } = useCunninghamTheme();
|
||||
const [isOpen, setIsOpen] = useState(false);
|
||||
const [buttonWidth, setButtonWidth] = useState<number | undefined>(undefined);
|
||||
const blockButtonRef = useRef<HTMLDivElement>(null);
|
||||
|
||||
useEffect(() => {
|
||||
// Mettre à jour la largeur uniquement côté client
|
||||
if (blockButtonRef.current) {
|
||||
setButtonWidth(blockButtonRef.current.clientWidth);
|
||||
}
|
||||
}, [isOpen]);
|
||||
|
||||
const onOpenChange = (isOpen: boolean) => {
|
||||
setIsOpen(isOpen);
|
||||
};
|
||||
@@ -76,7 +84,7 @@ export const DropdownMenu = ({
|
||||
>
|
||||
<Box
|
||||
$maxWidth="320px"
|
||||
$minWidth={`${blockButtonRef.current?.clientWidth}px`}
|
||||
$minWidth={buttonWidth ? `${buttonWidth}px` : undefined}
|
||||
role="menu"
|
||||
>
|
||||
{topMessage && (
|
||||
@@ -102,6 +110,7 @@ export const DropdownMenu = ({
|
||||
data-testid={option.testId}
|
||||
$direction="row"
|
||||
disabled={isDisabled}
|
||||
tabIndex={isDisabled ? -1 : 0}
|
||||
onClick={(event) => {
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
@@ -134,9 +143,15 @@ export const DropdownMenu = ({
|
||||
cursor: ${isDisabled ? 'not-allowed' : 'pointer'};
|
||||
user-select: none;
|
||||
|
||||
&:hover {
|
||||
&:hover,
|
||||
&:focus {
|
||||
background-color: var(--c--theme--colors--greyscale-050);
|
||||
}
|
||||
|
||||
&:focus-visible {
|
||||
outline: 2px solid #3e5de7;
|
||||
outline-offset: -2px;
|
||||
}
|
||||
`}
|
||||
>
|
||||
<Box
|
||||
|
||||
@@ -15,6 +15,7 @@ export const Icon = ({
|
||||
return (
|
||||
<Text
|
||||
{...textProps}
|
||||
aria-hidden="true"
|
||||
className={clsx('--docs--icon-bg', textProps.className, {
|
||||
'material-symbols': variant === 'filled',
|
||||
'material-symbols-outlined': variant === 'outlined',
|
||||
@@ -33,6 +34,7 @@ export const IconOptions = ({ isHorizontal, ...props }: IconOptionsProps) => {
|
||||
return (
|
||||
<Icon
|
||||
{...props}
|
||||
aria-hidden="true"
|
||||
iconName={isHorizontal ? 'more_horiz' : 'more_vert'}
|
||||
$css={css`
|
||||
user-select: none;
|
||||
|
||||
@@ -1,20 +1,17 @@
|
||||
import Lottie from 'react-lottie';
|
||||
import dynamic from 'next/dynamic';
|
||||
|
||||
const Lottie = dynamic(() => import('lottie-react'), { ssr: false });
|
||||
import searchingAnimation from '@/assets/lotties/searching';
|
||||
|
||||
export const Loader = () => {
|
||||
const LoaderOptions = {
|
||||
loop: true,
|
||||
autoplay: true,
|
||||
animationData: searchingAnimation,
|
||||
rendererSettings: {
|
||||
preserveAspectRatio: 'xMidYMid slice',
|
||||
},
|
||||
} as const;
|
||||
|
||||
export function Loader() {
|
||||
return (
|
||||
<div>
|
||||
<Lottie options={LoaderOptions} height={24} width={24} />
|
||||
<div role="status">
|
||||
<Lottie
|
||||
animationData={searchingAnimation}
|
||||
loop
|
||||
autoplay
|
||||
style={{ width: 24, height: 24 }}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ import {
|
||||
createContext,
|
||||
useCallback,
|
||||
useContext,
|
||||
useEffect,
|
||||
useState,
|
||||
} from 'react';
|
||||
import { createPortal } from 'react-dom';
|
||||
@@ -17,6 +18,8 @@ interface ToastItem {
|
||||
type: ToastType;
|
||||
icon?: string;
|
||||
duration?: number;
|
||||
actionLabel?: string;
|
||||
actionHref?: string;
|
||||
}
|
||||
|
||||
interface ToastContextType {
|
||||
@@ -25,6 +28,7 @@ interface ToastContextType {
|
||||
message: string,
|
||||
icon?: string,
|
||||
duration?: number,
|
||||
options?: { actionLabel?: string; actionHref?: string },
|
||||
) => void;
|
||||
}
|
||||
|
||||
@@ -44,11 +48,30 @@ interface ToastProviderProps {
|
||||
|
||||
export const ToastProvider = ({ children }: ToastProviderProps) => {
|
||||
const [toasts, setToasts] = useState<ToastItem[]>([]);
|
||||
const [isMounted, setIsMounted] = useState(false);
|
||||
|
||||
useEffect(() => {
|
||||
setIsMounted(true);
|
||||
}, []);
|
||||
|
||||
const showToast = useCallback(
|
||||
(type: ToastType, message: string, icon?: string, duration = 4000) => {
|
||||
(
|
||||
type: ToastType,
|
||||
message: string,
|
||||
icon?: string,
|
||||
duration = 4000,
|
||||
options?: { actionLabel?: string; actionHref?: string },
|
||||
) => {
|
||||
const id = Math.random().toString(36).substr(2, 9);
|
||||
const newToast: ToastItem = { id, message, type, icon, duration };
|
||||
const newToast: ToastItem = {
|
||||
id,
|
||||
message,
|
||||
type,
|
||||
icon,
|
||||
duration,
|
||||
actionLabel: options?.actionLabel,
|
||||
actionHref: options?.actionHref,
|
||||
};
|
||||
|
||||
setToasts((prev) => [newToast, ...prev]);
|
||||
},
|
||||
@@ -66,9 +89,12 @@ export const ToastProvider = ({ children }: ToastProviderProps) => {
|
||||
return (
|
||||
<ToastContext.Provider value={value}>
|
||||
{children}
|
||||
{typeof window !== 'undefined' &&
|
||||
{isMounted &&
|
||||
typeof document !== 'undefined' &&
|
||||
document.body &&
|
||||
createPortal(
|
||||
<Box
|
||||
aria-live="polite"
|
||||
$css={`
|
||||
position: fixed;
|
||||
top: 8px;
|
||||
@@ -94,6 +120,8 @@ export const ToastProvider = ({ children }: ToastProviderProps) => {
|
||||
type={toast.type}
|
||||
icon={toast.icon}
|
||||
duration={toast.duration}
|
||||
actionLabel={toast.actionLabel}
|
||||
actionHref={toast.actionHref}
|
||||
onClose={removeToast}
|
||||
/>
|
||||
))}
|
||||
|
||||
@@ -1,74 +1,64 @@
|
||||
import { CunninghamProvider } from '@openfun/cunningham-react';
|
||||
import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
|
||||
import { useRouter } from 'next/router';
|
||||
import dynamic from 'next/dynamic';
|
||||
import { useEffect } from 'react';
|
||||
|
||||
import { ToastProvider } from '@/components';
|
||||
import { useCunninghamTheme } from '@/cunningham';
|
||||
import { Auth, KEY_AUTH, setAuthUrl } from '@/features/auth';
|
||||
import { useResponsiveStore } from '@/stores/';
|
||||
import { useResponsiveStore } from '@/stores';
|
||||
|
||||
import { ConfigProvider } from './config/';
|
||||
import { ConfigProvider } from './config';
|
||||
|
||||
// Client-only providers
|
||||
const ToastProviderNoSSR = dynamic(
|
||||
() => import('@/components').then((mod) => ({ default: mod.ToastProvider })),
|
||||
{ ssr: false, loading: () => null },
|
||||
);
|
||||
|
||||
const CunninghamProviderNoSSR = dynamic(
|
||||
() =>
|
||||
import('@openfun/cunningham-react').then((mod) => ({
|
||||
default: mod.CunninghamProvider,
|
||||
})),
|
||||
{ ssr: false },
|
||||
);
|
||||
|
||||
/**
|
||||
* QueryClient:
|
||||
* - defaultOptions:
|
||||
* - staleTime:
|
||||
* - global cache duration - we decided 3 minutes
|
||||
* - It can be overridden to each query
|
||||
*/
|
||||
const defaultOptions = {
|
||||
queries: {
|
||||
staleTime: 1000 * 60 * 3,
|
||||
retry: 1,
|
||||
},
|
||||
};
|
||||
const queryClient = new QueryClient({
|
||||
defaultOptions,
|
||||
defaultOptions: {
|
||||
queries: { staleTime: 1000 * 60 * 3, retry: 1 },
|
||||
mutations: {
|
||||
onError: (error) => {
|
||||
if (
|
||||
error instanceof Error &&
|
||||
'status' in error &&
|
||||
error.status === 401
|
||||
) {
|
||||
void queryClient.resetQueries({ queryKey: [KEY_AUTH] });
|
||||
setAuthUrl();
|
||||
if (typeof window !== 'undefined') {
|
||||
window.location.href = '/401';
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
export function AppProvider({ children }: { children: React.ReactNode }) {
|
||||
const { theme } = useCunninghamTheme();
|
||||
const { replace } = useRouter();
|
||||
|
||||
const initializeResizeListener = useResponsiveStore(
|
||||
(state) => state.initializeResizeListener,
|
||||
);
|
||||
const theme = useCunninghamTheme((state) => state.theme);
|
||||
|
||||
useEffect(() => {
|
||||
return initializeResizeListener();
|
||||
}, [initializeResizeListener]);
|
||||
|
||||
useEffect(() => {
|
||||
queryClient.setDefaultOptions({
|
||||
...defaultOptions,
|
||||
mutations: {
|
||||
onError: (error) => {
|
||||
if (
|
||||
error instanceof Error &&
|
||||
'status' in error &&
|
||||
error.status === 401
|
||||
) {
|
||||
void queryClient.resetQueries({
|
||||
queryKey: [KEY_AUTH],
|
||||
});
|
||||
setAuthUrl();
|
||||
void replace(`/401`);
|
||||
}
|
||||
},
|
||||
},
|
||||
});
|
||||
}, [replace]);
|
||||
return useResponsiveStore.getState().initializeResizeListener();
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<QueryClientProvider client={queryClient}>
|
||||
<CunninghamProvider theme={theme}>
|
||||
<CunninghamProviderNoSSR theme={theme}>
|
||||
<ConfigProvider>
|
||||
<ToastProvider>
|
||||
<ToastProviderNoSSR>
|
||||
<Auth>{children}</Auth>
|
||||
</ToastProvider>
|
||||
</ToastProviderNoSSR>
|
||||
</ConfigProvider>
|
||||
</CunninghamProvider>
|
||||
</CunninghamProviderNoSSR>
|
||||
</QueryClientProvider>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -55,7 +55,7 @@ export const AttachmentList = ({
|
||||
>
|
||||
<Box
|
||||
$background="var(--c--theme--colors--greyscale-050)"
|
||||
$minWidth="200px"
|
||||
$width="200px"
|
||||
$direction="row"
|
||||
$gap="8px"
|
||||
$align="center"
|
||||
|
||||
@@ -8,6 +8,7 @@ import { Modal, ModalSize } from '@openfun/cunningham-react';
|
||||
import 'katex/dist/katex.min.css'; // `rehype-katex` does not import the CSS for you
|
||||
import { useRouter } from 'next/router';
|
||||
import { useCallback, useEffect, useRef, useState } from 'react';
|
||||
import type { ChangeEvent, FormEvent } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { MarkdownHooks } from 'react-markdown';
|
||||
import rehypeKatex from 'rehype-katex';
|
||||
@@ -16,7 +17,7 @@ import remarkGfm from 'remark-gfm';
|
||||
import remarkMath from 'remark-math';
|
||||
|
||||
import { APIError, errorCauses, fetchAPI } from '@/api';
|
||||
import { Box, Icon, Loader, StyledLink, Text } from '@/components';
|
||||
import { Box, Icon, Loader, Text } from '@/components';
|
||||
import { useUploadFile } from '@/features/attachments/hooks/useUploadFile';
|
||||
import { useChat } from '@/features/chat/api/useChat';
|
||||
import { getConversation } from '@/features/chat/api/useConversation';
|
||||
@@ -26,6 +27,8 @@ import {
|
||||
useLLMConfiguration,
|
||||
} from '@/features/chat/api/useLLMConfiguration';
|
||||
import { AttachmentList } from '@/features/chat/components/AttachmentList';
|
||||
import { ChatError } from '@/features/chat/components/ChatError';
|
||||
import { CodeBlock } from '@/features/chat/components/CodeBlock';
|
||||
import { FeedbackButtons } from '@/features/chat/components/FeedbackButtons';
|
||||
import { InputChat } from '@/features/chat/components/InputChat';
|
||||
import { SourceItemList } from '@/features/chat/components/SourceItemList';
|
||||
@@ -101,19 +104,14 @@ export const Chat = ({
|
||||
|
||||
if (modelToSelect) {
|
||||
setSelectedModel(modelToSelect);
|
||||
setSelectedModelHrid(modelToSelect.hrid);
|
||||
}
|
||||
}
|
||||
}, [llmConfig, selectedModel, selectedModelHrid]);
|
||||
|
||||
// Update store when model selection changes
|
||||
useEffect(() => {
|
||||
if (selectedModel?.hrid !== selectedModelHrid) {
|
||||
setSelectedModelHrid(selectedModel?.hrid || null);
|
||||
}
|
||||
}, [selectedModel, selectedModelHrid, setSelectedModelHrid]);
|
||||
}, [llmConfig, selectedModel, selectedModelHrid, setSelectedModelHrid]);
|
||||
|
||||
const handleModelSelect = (model: LLMModel) => {
|
||||
setSelectedModel(model);
|
||||
setSelectedModelHrid(model.hrid);
|
||||
};
|
||||
|
||||
const router = useRouter();
|
||||
@@ -153,17 +151,26 @@ export const Chat = ({
|
||||
const [initialConversationMessages, setInitialConversationMessages] =
|
||||
useState<Message[] | undefined>(undefined);
|
||||
const [pendingFirstMessage, setPendingFirstMessage] = useState<{
|
||||
event: React.FormEvent<HTMLFormElement>;
|
||||
event: FormEvent<HTMLFormElement>;
|
||||
attachments?: Attachment[];
|
||||
forceWebSearch?: boolean;
|
||||
} | null>(null);
|
||||
const [shouldAutoSubmit, setShouldAutoSubmit] = useState(false);
|
||||
const [shouldRetry, setShouldRetry] = useState(false);
|
||||
const retryOriginalInputRef = useRef<string>('');
|
||||
const retryOriginalFilesRef = useRef<FileList | null>(null);
|
||||
const [hasInitialized, setHasInitialized] = useState(false);
|
||||
const [streamingMessageHeight, setStreamingMessageHeight] = useState<
|
||||
number | null
|
||||
>(null);
|
||||
const lastUserMessageIdRef = useRef<string | null>(null);
|
||||
const hasScrolledToBottomOnLoadRef = useRef(false);
|
||||
const lastSubmissionRef = useRef<{
|
||||
input: string;
|
||||
files: FileList | null;
|
||||
event: FormEvent<HTMLFormElement>;
|
||||
options?: Record<string, unknown>;
|
||||
} | null>(null);
|
||||
|
||||
const { mutate: createChatConversation } = useCreateChatConversation();
|
||||
|
||||
@@ -212,6 +219,7 @@ export const Chat = ({
|
||||
handleInputChange,
|
||||
status,
|
||||
stop: stopChat,
|
||||
setMessages,
|
||||
} = useChat({
|
||||
id: conversationId,
|
||||
initialMessages: initialConversationMessages,
|
||||
@@ -244,10 +252,33 @@ export const Chat = ({
|
||||
void stopGeneration();
|
||||
};
|
||||
|
||||
const handleSubmitWrapper = (event: React.FormEvent<HTMLFormElement>) => {
|
||||
const handleSubmitWrapper = (event: FormEvent<HTMLFormElement>) => {
|
||||
void handleSubmit(event);
|
||||
};
|
||||
|
||||
const handleRetry = () => {
|
||||
if (!lastSubmissionRef.current || !setMessages) {
|
||||
return;
|
||||
}
|
||||
|
||||
const { input: lastInput, files: lastFiles } = lastSubmissionRef.current;
|
||||
|
||||
const lastAssistantIndex = messages.findLastIndex(
|
||||
(msg) => msg.role === 'assistant',
|
||||
);
|
||||
if (lastAssistantIndex !== -1) {
|
||||
setMessages(messages.filter((_, index) => index !== lastAssistantIndex));
|
||||
}
|
||||
|
||||
retryOriginalInputRef.current = input;
|
||||
retryOriginalFilesRef.current = files;
|
||||
handleInputChange({
|
||||
target: { value: lastInput },
|
||||
} as ChangeEvent<HTMLTextAreaElement>);
|
||||
setFiles(lastFiles);
|
||||
setShouldRetry(true);
|
||||
};
|
||||
|
||||
// Précharger les métadonnées des sources dès que les messages arrivent
|
||||
useEffect(() => {
|
||||
messages.forEach((message) => {
|
||||
@@ -260,7 +291,8 @@ export const Chat = ({
|
||||
});
|
||||
}
|
||||
});
|
||||
}, [messages, prefetchMetadata]);
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [messages]);
|
||||
|
||||
const openSources = (messageId: string) => {
|
||||
if (isSourceOpen === messageId) {
|
||||
@@ -303,16 +335,23 @@ export const Chat = ({
|
||||
|
||||
const availableHeight = containerHeight - userMessageHeight - 38;
|
||||
|
||||
if (streamingMessageHeight !== availableHeight) {
|
||||
setStreamingMessageHeight(availableHeight);
|
||||
}
|
||||
setStreamingMessageHeight((prev) => {
|
||||
if (prev === null || Math.abs(prev - availableHeight) > 10) {
|
||||
return availableHeight;
|
||||
}
|
||||
return prev;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}, [messages, streamingMessageHeight]);
|
||||
}, [messages]);
|
||||
|
||||
// Détecter l'arrivée d'un nouveau message user et retirer la hauteur de l'ancien
|
||||
useEffect(() => {
|
||||
if (status === 'streaming') {
|
||||
return;
|
||||
}
|
||||
|
||||
const userMessages = messages.filter((msg) => msg.role === 'user');
|
||||
const lastUserMessage = userMessages[userMessages.length - 1];
|
||||
|
||||
@@ -325,14 +364,14 @@ export const Chat = ({
|
||||
}
|
||||
lastUserMessageIdRef.current = lastUserMessage.id;
|
||||
}
|
||||
}, [messages]);
|
||||
}, [messages, status]);
|
||||
|
||||
// Calculer la hauteur pendant submitted/streaming
|
||||
useEffect(() => {
|
||||
if (status === 'submitted' || status === 'streaming') {
|
||||
calculateStreamingHeight();
|
||||
}
|
||||
}, [status, calculateStreamingHeight]);
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [status]);
|
||||
|
||||
// Scroller vers la question au moment du submit
|
||||
useEffect(() => {
|
||||
@@ -352,7 +391,8 @@ export const Chat = ({
|
||||
|
||||
messageElement?.scrollIntoView({ block: 'start', behavior: 'smooth' });
|
||||
});
|
||||
}, [status, messages]);
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [status]);
|
||||
|
||||
// Synchronize conversationId state with prop when it changes (e.g., after navigation)
|
||||
useEffect(() => {
|
||||
@@ -361,10 +401,11 @@ export const Chat = ({
|
||||
if (initialConversationId !== conversationId) {
|
||||
handleInputChange({
|
||||
target: { value: '' },
|
||||
} as React.ChangeEvent<HTMLTextAreaElement>);
|
||||
} as ChangeEvent<HTMLTextAreaElement>);
|
||||
setHasInitialized(false); // Réinitialiser pour permettre le scroll au prochain chargement
|
||||
}
|
||||
}, [initialConversationId, conversationId, handleInputChange]);
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [initialConversationId, conversationId]);
|
||||
|
||||
// On mount, if there is pending input/files, initialize state and set flag
|
||||
useEffect(() => {
|
||||
@@ -375,7 +416,7 @@ export const Chat = ({
|
||||
if (pendingInput) {
|
||||
const syntheticEvent = {
|
||||
target: { value: pendingInput },
|
||||
} as React.ChangeEvent<HTMLInputElement>;
|
||||
} as ChangeEvent<HTMLInputElement>;
|
||||
handleInputChange(syntheticEvent);
|
||||
}
|
||||
if (pendingFiles) {
|
||||
@@ -397,13 +438,32 @@ export const Chat = ({
|
||||
const syntheticFormEvent = {
|
||||
preventDefault: () => {},
|
||||
target: form,
|
||||
} as unknown as React.FormEvent<HTMLFormElement>;
|
||||
} as unknown as FormEvent<HTMLFormElement>;
|
||||
void handleSubmit(syntheticFormEvent);
|
||||
setShouldAutoSubmit(false);
|
||||
}
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [shouldAutoSubmit, input, files]);
|
||||
|
||||
useEffect(() => {
|
||||
if (
|
||||
shouldRetry &&
|
||||
lastSubmissionRef.current &&
|
||||
input === lastSubmissionRef.current.input
|
||||
) {
|
||||
const { event } = lastSubmissionRef.current;
|
||||
|
||||
void handleSubmit(event);
|
||||
handleInputChange({
|
||||
target: { value: retryOriginalInputRef.current },
|
||||
} as ChangeEvent<HTMLTextAreaElement>);
|
||||
setFiles(retryOriginalFilesRef.current);
|
||||
|
||||
setShouldRetry(false);
|
||||
}
|
||||
// eslint-disable-next-line react-hooks/exhaustive-deps
|
||||
}, [shouldRetry, input, files]);
|
||||
|
||||
// Fetch initial conversation messages if initialConversationId is provided and no pending input
|
||||
useEffect(() => {
|
||||
hasScrolledToBottomOnLoadRef.current = false; // Réinitialiser au début du chargement
|
||||
@@ -457,7 +517,7 @@ export const Chat = ({
|
||||
}, [hasInitialized, messages.length]);
|
||||
|
||||
// Custom handleSubmit to include attachments and handle chat creation
|
||||
const handleSubmit = async (event: React.FormEvent<HTMLFormElement>) => {
|
||||
const handleSubmit = async (event: FormEvent<HTMLFormElement>) => {
|
||||
event.preventDefault();
|
||||
|
||||
// Upload files to server and get URLs
|
||||
@@ -536,6 +596,13 @@ export const Chat = ({
|
||||
options.experimental_attachments = attachments;
|
||||
}
|
||||
|
||||
lastSubmissionRef.current = {
|
||||
input,
|
||||
files,
|
||||
event,
|
||||
options: Object.keys(options).length > 0 ? options : undefined,
|
||||
};
|
||||
|
||||
if (Object.keys(options).length > 0) {
|
||||
baseHandleSubmit(event, options);
|
||||
} else {
|
||||
@@ -603,6 +670,7 @@ export const Chat = ({
|
||||
display: flex;
|
||||
width: 100%;
|
||||
margin: auto;
|
||||
margin-bottom: ${isLastAssistantMessageInConversation ? '30px' : '0px'};
|
||||
color: var(--c--theme--colors--greyscale-850);
|
||||
padding-left: 12px;
|
||||
padding-right: 12px;
|
||||
@@ -641,34 +709,65 @@ export const Chat = ({
|
||||
{message.content && (
|
||||
<Box
|
||||
className="mainContent-chat"
|
||||
data-testid={
|
||||
message.role === 'assistant'
|
||||
? 'assistant-message-content'
|
||||
: undefined
|
||||
}
|
||||
$padding={{ all: 'xxs' }}
|
||||
>
|
||||
<MarkdownHooks
|
||||
remarkPlugins={[remarkGfm, remarkMath]}
|
||||
rehypePlugins={[
|
||||
[
|
||||
rehypePrettyCode,
|
||||
{
|
||||
theme: 'github-dark-dimmed',
|
||||
},
|
||||
],
|
||||
rehypeKatex,
|
||||
]}
|
||||
components={{
|
||||
// Custom components for Markdown rendering
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
p: ({ node, ...props }) => (
|
||||
<Text
|
||||
$css="display: block"
|
||||
$theme="greyscale"
|
||||
$variation="850"
|
||||
{...props}
|
||||
/>
|
||||
),
|
||||
}}
|
||||
>
|
||||
{message.content}
|
||||
</MarkdownHooks>
|
||||
<p className="sr-only">
|
||||
{message.role === 'user'
|
||||
? t('You said: ')
|
||||
: t('Assistant IA replied: ')}
|
||||
</p>
|
||||
{message.role === 'user' ? (
|
||||
<Text
|
||||
as="p"
|
||||
$css="white-space: pre-wrap; display: block;"
|
||||
$theme="greyscale"
|
||||
$variation="850"
|
||||
>
|
||||
{message.content}
|
||||
</Text>
|
||||
) : (
|
||||
<MarkdownHooks
|
||||
remarkPlugins={[remarkGfm, remarkMath]}
|
||||
rehypePlugins={[
|
||||
[
|
||||
rehypePrettyCode,
|
||||
{
|
||||
theme: 'github-dark-dimmed',
|
||||
},
|
||||
],
|
||||
rehypeKatex,
|
||||
]}
|
||||
components={{
|
||||
// Custom components for Markdown rendering
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
p: ({ node, ...props }) => (
|
||||
<Text
|
||||
as="p"
|
||||
$css="display: block"
|
||||
$theme="greyscale"
|
||||
$variation="850"
|
||||
{...props}
|
||||
/>
|
||||
),
|
||||
a: ({ children, ...props }) => (
|
||||
<a target="_blank" {...props}>
|
||||
{children}
|
||||
</a>
|
||||
),
|
||||
// eslint-disable-next-line @typescript-eslint/no-unused-vars
|
||||
pre: ({ node, children, ...props }) => (
|
||||
<CodeBlock {...props}>{children}</CodeBlock>
|
||||
),
|
||||
}}
|
||||
>
|
||||
{message.content}
|
||||
</MarkdownHooks>
|
||||
)}
|
||||
</Box>
|
||||
)}
|
||||
|
||||
@@ -696,7 +795,23 @@ export const Chat = ({
|
||||
>
|
||||
<Loader />
|
||||
<Text $variation="600" $size="md">
|
||||
{t('Search...')}
|
||||
{(() => {
|
||||
const toolInvocation = message.parts?.find(
|
||||
(part) =>
|
||||
part.type === 'tool-invocation' &&
|
||||
part.toolInvocation.toolName !==
|
||||
'document_parsing',
|
||||
);
|
||||
if (
|
||||
toolInvocation?.type ===
|
||||
'tool-invocation' &&
|
||||
toolInvocation.toolInvocation.toolName ===
|
||||
'summarize'
|
||||
) {
|
||||
return t('Summarizing...');
|
||||
}
|
||||
return t('Search...');
|
||||
})()}
|
||||
</Text>
|
||||
</Box>
|
||||
)}
|
||||
@@ -885,28 +1000,11 @@ export const Chat = ({
|
||||
</Text>
|
||||
</Box>
|
||||
) : null}
|
||||
{status === 'error' && (
|
||||
<Box
|
||||
$direction={isMobile ? 'column' : 'row'}
|
||||
$align="center"
|
||||
$gap="6px"
|
||||
$width="100%"
|
||||
$maxWidth="750px"
|
||||
$margin={{ all: 'auto', top: 'base', bottom: 'md' }}
|
||||
$padding={{ left: '13px' }}
|
||||
>
|
||||
<Text>{t('Sorry, an error occurred. Please try again.')}</Text>
|
||||
<StyledLink
|
||||
href="/"
|
||||
rel="noopener noreferrer"
|
||||
$css={`
|
||||
color: var(--c--theme--colors--greyscale-900);
|
||||
text-decoration: underline;
|
||||
`}
|
||||
>
|
||||
{t('Start a new conversation.')}
|
||||
</StyledLink>
|
||||
</Box>
|
||||
{status === 'error' && !isUploadingFiles && (
|
||||
<ChatError
|
||||
hasLastSubmission={!!lastSubmissionRef.current}
|
||||
onRetry={handleRetry}
|
||||
/>
|
||||
)}
|
||||
</Box>
|
||||
<Box
|
||||
@@ -914,6 +1012,7 @@ export const Chat = ({
|
||||
position: relative;
|
||||
bottom: ${isMobile ? '8px' : '20px'};
|
||||
margin: auto;
|
||||
background-color: white;
|
||||
z-index: 1000;
|
||||
`}
|
||||
$gap="6px"
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
import { Button } from '@openfun/cunningham-react';
|
||||
import { useRouter } from 'next/router';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import { Box, Icon, Text } from '@/components';
|
||||
|
||||
interface ChatErrorProps {
|
||||
hasLastSubmission: boolean;
|
||||
onRetry: () => void;
|
||||
}
|
||||
|
||||
export const ChatError = ({ hasLastSubmission, onRetry }: ChatErrorProps) => {
|
||||
const { t } = useTranslation();
|
||||
const router = useRouter();
|
||||
|
||||
return (
|
||||
<Box
|
||||
$direction="column"
|
||||
$gap="6px"
|
||||
$width="100%"
|
||||
$maxWidth="750px"
|
||||
$margin={{ all: 'auto', top: 'base', bottom: 'md' }}
|
||||
$padding={{ left: '13px' }}
|
||||
>
|
||||
<Text $variation="550" $theme="greyscale">
|
||||
{t('Sorry, an error occurred. Please try again.')}
|
||||
</Text>
|
||||
<Box
|
||||
$direction="row"
|
||||
$gap="6px"
|
||||
$align="center"
|
||||
$margin={{ top: '10px' }}
|
||||
>
|
||||
{hasLastSubmission ? (
|
||||
<Button
|
||||
size="small"
|
||||
color="tertiary"
|
||||
onClick={onRetry}
|
||||
className="retry-button"
|
||||
style={{
|
||||
color: 'var(--c--theme--colors--greyscale-550)',
|
||||
borderColor: 'var(--c--theme--colors--greyscale-300)',
|
||||
}}
|
||||
icon={
|
||||
<svg
|
||||
xmlns="http://www.w3.org/2000/svg"
|
||||
width="11"
|
||||
height="15"
|
||||
viewBox="0 0 11 15"
|
||||
fill="none"
|
||||
>
|
||||
<path
|
||||
d="M0.733333 10.0333C0.488889 9.61111 0.305556 9.17778 0.183333 8.73333C0.0611111 8.28889 0 7.83333 0 7.36667C0 5.87778 0.516667 4.61111 1.55 3.56667C2.58333 2.52222 3.84444 2 5.33333 2H5.45L4.38333 0.933333L5.31667 0L7.98333 2.66667L5.31667 5.33333L4.38333 4.4L5.45 3.33333H5.33333C4.22222 3.33333 3.27778 3.725 2.5 4.50833C1.72222 5.29167 1.33333 6.24444 1.33333 7.36667C1.33333 7.65556 1.36667 7.93889 1.43333 8.21667C1.5 8.49444 1.6 8.76667 1.73333 9.03333L0.733333 10.0333ZM5.35 14.6667L2.68333 12L5.35 9.33333L6.28333 10.2667L5.21667 11.3333H5.33333C6.44444 11.3333 7.38889 10.9417 8.16667 10.1583C8.94444 9.375 9.33333 8.42222 9.33333 7.3C9.33333 7.01111 9.3 6.72778 9.23333 6.45C9.16667 6.17222 9.06667 5.9 8.93333 5.63333L9.93333 4.63333C10.1778 5.05556 10.3611 5.48889 10.4833 5.93333C10.6056 6.37778 10.6667 6.83333 10.6667 7.3C10.6667 8.78889 10.15 10.0556 9.11667 11.1C8.08333 12.1444 6.82222 12.6667 5.33333 12.6667H5.21667L6.28333 13.7333L5.35 14.6667Z"
|
||||
fill="currentColor"
|
||||
/>
|
||||
</svg>
|
||||
}
|
||||
>
|
||||
{t('Retry')}
|
||||
</Button>
|
||||
) : (
|
||||
<Button
|
||||
size="small"
|
||||
color="tertiary"
|
||||
style={{
|
||||
color: 'var(--c--theme--colors--greyscale-550)',
|
||||
borderColor: 'var(--c--theme--colors--greyscale-300)',
|
||||
}}
|
||||
onClick={() => {
|
||||
void router.push('/');
|
||||
}}
|
||||
icon={<Icon iconName="add" $color="greyscale" />}
|
||||
>
|
||||
{t('Start a new conversation')}
|
||||
</Button>
|
||||
)}
|
||||
</Box>
|
||||
</Box>
|
||||
);
|
||||
};
|
||||
@@ -0,0 +1,80 @@
|
||||
import { useRef } from 'react';
|
||||
import type { ReactNode } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import { Box, Icon, Text } from '@/components';
|
||||
import { useClipboard } from '@/hook';
|
||||
|
||||
interface CopyCodeButtonProps {
|
||||
onCopy: () => void;
|
||||
}
|
||||
|
||||
const CopyCodeButton = ({ onCopy }: CopyCodeButtonProps) => {
|
||||
const { t } = useTranslation();
|
||||
|
||||
return (
|
||||
<Box
|
||||
as="button"
|
||||
onClick={onCopy}
|
||||
$css={`
|
||||
position: absolute;
|
||||
top: 8px;
|
||||
right: 8px;
|
||||
padding: 6px 10px;
|
||||
background: rgba(0, 0, 0, 0.03);
|
||||
border: 1px solid rgba(255, 255, 255, 0.15);
|
||||
border-radius: 4px;
|
||||
cursor: pointer;
|
||||
font-size: 12px;
|
||||
font-weight: 500;
|
||||
color: #fff;
|
||||
display: flex;
|
||||
flex-direction: row;
|
||||
align-items: center;
|
||||
gap: 4px;
|
||||
z-index: 10;
|
||||
transition: all 0.2s;
|
||||
box-shadow: 0 1px 2px rgba(0, 0, 0, 0.1);
|
||||
width: fit-content;
|
||||
&:hover {
|
||||
background:rgba(255, 255, 255, 0.1);
|
||||
border: 1px solid rgba(255, 255, 255, 0.20);
|
||||
}
|
||||
`}
|
||||
>
|
||||
<Icon
|
||||
iconName="content_copy"
|
||||
$size="14px"
|
||||
$theme="greyscale"
|
||||
$variation="200"
|
||||
/>
|
||||
<Text $size="xs" $theme="greyscale" $variation="200">
|
||||
{t('Copy code')}
|
||||
</Text>
|
||||
</Box>
|
||||
);
|
||||
};
|
||||
|
||||
interface CodeBlockProps {
|
||||
children: ReactNode;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
export const CodeBlock = ({ children, ...props }: CodeBlockProps) => {
|
||||
const preRef = useRef<HTMLPreElement>(null);
|
||||
const copyToClipboard = useClipboard();
|
||||
|
||||
const handleCopy = () => {
|
||||
const code = preRef.current?.querySelector('code');
|
||||
copyToClipboard(code?.textContent || '');
|
||||
};
|
||||
|
||||
return (
|
||||
<>
|
||||
<CopyCodeButton onCopy={handleCopy} />
|
||||
<Box ref={preRef} $position="relative" as="pre" {...props}>
|
||||
{children}
|
||||
</Box>
|
||||
</>
|
||||
);
|
||||
};
|
||||
@@ -1,8 +1,9 @@
|
||||
import { Button } from '@openfun/cunningham-react';
|
||||
import React, { useEffect, useRef, useState } from 'react';
|
||||
import React, { useCallback, useEffect, useRef, useState } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
|
||||
import { Box, Icon, Text } from '@/components';
|
||||
import { useToast } from '@/components/ToastProvider';
|
||||
import { FeatureFlagState, useConfig } from '@/core';
|
||||
import { LLMModel } from '@/features/chat/api/useLLMConfiguration';
|
||||
import { useAnalytics } from '@/libs';
|
||||
@@ -51,10 +52,10 @@ export const InputChat = ({
|
||||
isUploadingFiles = false,
|
||||
}: InputChatProps) => {
|
||||
const { t } = useTranslation();
|
||||
const { showToast } = useToast();
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
const textareaRef = useRef<HTMLTextAreaElement>(null);
|
||||
const [isDragActive, setIsDragActive] = useState(false);
|
||||
const [isDragRejected, setIsDragRejected] = useState(false);
|
||||
const { isDesktop, isMobile } = useResponsiveStore();
|
||||
const [currentSuggestionIndex, setCurrentSuggestionIndex] = useState(0);
|
||||
const { data: conf } = useConfig();
|
||||
@@ -63,6 +64,29 @@ export const InputChat = ({
|
||||
const [webSearchEnabled, setWebSearchEnabled] = useState(false);
|
||||
const [isResetting, setIsResetting] = useState(false);
|
||||
|
||||
const isFileAccepted = useCallback(
|
||||
(file: File): boolean => {
|
||||
const acceptedConfig = conf?.chat_upload_accept;
|
||||
if (!acceptedConfig) {
|
||||
return true;
|
||||
}
|
||||
const acceptedTypes = acceptedConfig
|
||||
.split(',')
|
||||
.map((type) => type.trim());
|
||||
return acceptedTypes.some((acceptedType) => {
|
||||
if (acceptedType.startsWith('.')) {
|
||||
return file.name.toLowerCase().endsWith(acceptedType.toLowerCase());
|
||||
}
|
||||
if (acceptedType.endsWith('/*')) {
|
||||
const baseType = acceptedType.slice(0, -2);
|
||||
return file.type.startsWith(baseType);
|
||||
}
|
||||
return file.type === acceptedType;
|
||||
});
|
||||
},
|
||||
[conf?.chat_upload_accept],
|
||||
);
|
||||
|
||||
const suggestions = [
|
||||
t('Ask a question'),
|
||||
t('Turn this list into bullet points'),
|
||||
@@ -70,6 +94,20 @@ export const InputChat = ({
|
||||
t('Find recent news about...'),
|
||||
];
|
||||
|
||||
const showToastError = useCallback(() => {
|
||||
showToast(
|
||||
'error',
|
||||
`${t('File type not supported')}`,
|
||||
undefined,
|
||||
undefined,
|
||||
{
|
||||
actionLabel: t('Know more'),
|
||||
actionHref:
|
||||
'https://docs.numerique.gouv.fr/docs/060b7b70-15aa-4d9a-86f5-2d31c3d693d5/',
|
||||
},
|
||||
);
|
||||
}, [showToast, t]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!conf?.FEATURE_FLAGS) {
|
||||
setWebSearchEnabled(false);
|
||||
@@ -136,39 +174,10 @@ export const InputChat = ({
|
||||
return;
|
||||
}
|
||||
|
||||
const isFileAccepted = (file: File): boolean => {
|
||||
if (!conf?.chat_upload_accept) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const acceptedTypes = conf.chat_upload_accept
|
||||
.split(',')
|
||||
.map((type) => type.trim());
|
||||
|
||||
return acceptedTypes.some((acceptedType) => {
|
||||
// Extension management
|
||||
if (acceptedType.startsWith('.')) {
|
||||
return file.name.toLowerCase().endsWith(acceptedType.toLowerCase());
|
||||
}
|
||||
// Wildcard MIME type management (ex: image/*)
|
||||
if (acceptedType.endsWith('/*')) {
|
||||
const baseType = acceptedType.slice(0, -2);
|
||||
return file.type.startsWith(baseType);
|
||||
}
|
||||
// Exact MIME type management
|
||||
return file.type === acceptedType;
|
||||
});
|
||||
};
|
||||
|
||||
const areAllFilesAccepted = (fileList: FileList): boolean => {
|
||||
return Array.from(fileList).every((file) => isFileAccepted(file));
|
||||
};
|
||||
|
||||
const handleDragEnter = (e: DragEvent) => {
|
||||
e.preventDefault();
|
||||
if (e.dataTransfer?.types.includes('Files')) {
|
||||
setIsDragActive(true);
|
||||
setIsDragRejected(false);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -177,7 +186,6 @@ export const InputChat = ({
|
||||
// Only hide when leaving the window completely
|
||||
if (!e.relatedTarget) {
|
||||
setIsDragActive(false);
|
||||
setIsDragRejected(false);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -187,7 +195,7 @@ export const InputChat = ({
|
||||
// Check for rejected files during drag over (does not work on Safari)
|
||||
if (e.dataTransfer?.items) {
|
||||
const items = Array.from(e.dataTransfer.items);
|
||||
const hasInvalidFile = items.some((item) => {
|
||||
items.some((item) => {
|
||||
if (item.kind === 'file') {
|
||||
// Check file type
|
||||
const type = item.type;
|
||||
@@ -196,15 +204,12 @@ export const InputChat = ({
|
||||
}
|
||||
return false;
|
||||
});
|
||||
|
||||
setIsDragRejected(hasInvalidFile);
|
||||
}
|
||||
};
|
||||
|
||||
const handleDrop = (e: DragEvent) => {
|
||||
e.preventDefault();
|
||||
setIsDragActive(false);
|
||||
setIsDragRejected(false);
|
||||
|
||||
if (!fileUploadEnabled) {
|
||||
return;
|
||||
@@ -212,15 +217,22 @@ export const InputChat = ({
|
||||
|
||||
const droppedFiles = e.dataTransfer?.files;
|
||||
if (droppedFiles && droppedFiles.length > 0) {
|
||||
// Check if all files are accepted
|
||||
if (!areAllFilesAccepted(droppedFiles)) {
|
||||
// Display rejection for 2 seconds (mandatory for Safari)
|
||||
setIsDragActive(true);
|
||||
setIsDragRejected(true);
|
||||
setTimeout(() => {
|
||||
setIsDragActive(false);
|
||||
setIsDragRejected(false);
|
||||
}, 2000);
|
||||
const acceptedFiles: File[] = [];
|
||||
const rejectedFiles: string[] = [];
|
||||
|
||||
Array.from(droppedFiles).forEach((file) => {
|
||||
if (isFileAccepted(file)) {
|
||||
acceptedFiles.push(file);
|
||||
} else {
|
||||
rejectedFiles.push(file.name);
|
||||
}
|
||||
});
|
||||
|
||||
if (rejectedFiles.length > 0) {
|
||||
showToastError();
|
||||
}
|
||||
|
||||
if (acceptedFiles.length === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -229,7 +241,7 @@ export const InputChat = ({
|
||||
if (prev) {
|
||||
Array.from(prev).forEach((f) => dt.items.add(f));
|
||||
}
|
||||
Array.from(droppedFiles).forEach((f) => {
|
||||
acceptedFiles.forEach((f) => {
|
||||
if (
|
||||
!Array.from(prev || []).some(
|
||||
(pf) =>
|
||||
@@ -257,7 +269,13 @@ export const InputChat = ({
|
||||
window.removeEventListener('dragover', handleDragOver);
|
||||
window.removeEventListener('drop', handleDrop);
|
||||
};
|
||||
}, [fileUploadEnabled, setFiles, conf?.chat_upload_accept]);
|
||||
}, [
|
||||
fileUploadEnabled,
|
||||
setFiles,
|
||||
showToastError,
|
||||
conf?.chat_upload_accept,
|
||||
isFileAccepted,
|
||||
]);
|
||||
|
||||
const isInputDisabled = status !== 'ready' || isUploadingFiles;
|
||||
|
||||
@@ -282,7 +300,6 @@ export const InputChat = ({
|
||||
$css={`
|
||||
display: block;
|
||||
position: relative;
|
||||
opacity: ${status === 'error' ? '0.5' : '1'};
|
||||
margin: auto;
|
||||
width: 100%;
|
||||
padding: ${isDesktop ? '0' : '0 10px'};
|
||||
@@ -290,26 +307,22 @@ export const InputChat = ({
|
||||
`}
|
||||
>
|
||||
{/* Bouton de scroll vers le bas */}
|
||||
{messagesLength > 1 &&
|
||||
status !== 'streaming' &&
|
||||
status !== 'submitted' &&
|
||||
containerRef &&
|
||||
onScrollToBottom && (
|
||||
<Box
|
||||
$css={`
|
||||
{messagesLength > 1 && containerRef && onScrollToBottom && (
|
||||
<Box
|
||||
$css={`
|
||||
position: relative;
|
||||
height: 0;
|
||||
width: 100%;
|
||||
margin: auto;
|
||||
max-width: 750px;
|
||||
`}
|
||||
>
|
||||
<ScrollDown
|
||||
onClick={onScrollToBottom}
|
||||
containerRef={containerRef}
|
||||
/>
|
||||
</Box>
|
||||
)}
|
||||
>
|
||||
<ScrollDown
|
||||
onClick={onScrollToBottom}
|
||||
containerRef={containerRef}
|
||||
/>
|
||||
</Box>
|
||||
)}
|
||||
{/* Message de bienvenue */}
|
||||
{messagesLength === 0 && (
|
||||
<Box
|
||||
@@ -365,52 +378,27 @@ export const InputChat = ({
|
||||
top: -1px; left: -1px;
|
||||
border-radius: 12px;
|
||||
z-index: 1001;
|
||||
background-color: ${isDragRejected ? '#FFE8E8' : '#EDF0FF'};
|
||||
background-color: #EDF0FF;
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
outline: 2px solid ${isDragRejected ? '#FF6B6B' : '#90A7FF'};
|
||||
box-shadow: 0 0 64px 0 ${isDragRejected ? 'rgba(255, 107, 107, 0.25)' : 'rgba(62, 93, 231, 0.25)'};
|
||||
outline: 2px solid #90A7FF;
|
||||
box-shadow: 0 0 64px 0 rgba(62, 93, 231, 0.25);
|
||||
`}
|
||||
>
|
||||
{isDragRejected ? (
|
||||
<>
|
||||
<Text $css="font-size: 48px;">🚫</Text>
|
||||
<Box>
|
||||
<Text $weight="700" $color="#C92A2A">
|
||||
{t('File type not supported (yet)')}
|
||||
</Text>
|
||||
<Text $weight="400" $color="#C92A2A">
|
||||
{t(
|
||||
'We currently support only specific file types...',
|
||||
)}
|
||||
</Text>
|
||||
<Text $weight="400" $color="#C92A2A">
|
||||
{t(
|
||||
'Use the "{{attach_file_btn}}" button to have a better view.',
|
||||
{ attach_file_btn: t('Add attach file') },
|
||||
)}
|
||||
</Text>
|
||||
</Box>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<FilesIcon />
|
||||
<Box>
|
||||
<Text $weight="700" $color="#223E9E">
|
||||
{t('Add file')}
|
||||
</Text>
|
||||
<Text $weight="400" $color="#223E9E">
|
||||
{t(
|
||||
'To add a file to the conversation, drop it here.',
|
||||
)}
|
||||
</Text>
|
||||
</Box>
|
||||
</>
|
||||
)}
|
||||
<FilesIcon />
|
||||
<Box>
|
||||
<Text $weight="700" $color="#223E9E">
|
||||
{t('Add file')}
|
||||
</Text>
|
||||
<Text $weight="400" $color="#223E9E">
|
||||
{t('To add a file to the conversation, drop it here.')}
|
||||
</Text>
|
||||
</Box>
|
||||
</Box>
|
||||
)}
|
||||
<textarea
|
||||
ref={textareaRef}
|
||||
aria-label={t('Enter your message or a question')}
|
||||
value={input ?? ''}
|
||||
name="inputchat-textarea"
|
||||
onChange={(e) => {
|
||||
@@ -430,6 +418,7 @@ export const InputChat = ({
|
||||
fontSize: '1rem',
|
||||
border: 'none',
|
||||
resize: 'none',
|
||||
opacity: status === 'error' ? '0.5' : '1',
|
||||
fontFamily: 'inherit',
|
||||
minHeight: '64px',
|
||||
maxHeight: '200px',
|
||||
@@ -506,12 +495,33 @@ export const InputChat = ({
|
||||
if (!fileList) {
|
||||
return;
|
||||
}
|
||||
|
||||
const acceptedFiles: File[] = [];
|
||||
const rejectedFiles: string[] = [];
|
||||
|
||||
Array.from(fileList).forEach((file) => {
|
||||
if (isFileAccepted(file)) {
|
||||
acceptedFiles.push(file);
|
||||
} else {
|
||||
rejectedFiles.push(file.name);
|
||||
}
|
||||
});
|
||||
|
||||
if (rejectedFiles.length > 0) {
|
||||
showToastError();
|
||||
}
|
||||
|
||||
if (acceptedFiles.length === 0) {
|
||||
e.target.value = '';
|
||||
return;
|
||||
}
|
||||
|
||||
setFiles((prev) => {
|
||||
const dt = new DataTransfer();
|
||||
if (prev) {
|
||||
Array.from(prev).forEach((f: File) => dt.items.add(f));
|
||||
}
|
||||
Array.from(fileList).forEach((f: File) => {
|
||||
acceptedFiles.forEach((f: File) => {
|
||||
if (
|
||||
!Array.from(prev || []).some(
|
||||
(pf) =>
|
||||
@@ -525,6 +535,8 @@ export const InputChat = ({
|
||||
});
|
||||
return dt.files;
|
||||
});
|
||||
|
||||
e.target.value = '';
|
||||
}}
|
||||
/>
|
||||
{/*Aperçu des fichiers*/}
|
||||
@@ -557,6 +569,9 @@ export const InputChat = ({
|
||||
$gap="sm"
|
||||
$padding={{ bottom: 'base' }}
|
||||
$align="space-between"
|
||||
$css={`
|
||||
opacity: ${status === 'error' ? '0.5' : '1'};
|
||||
`}
|
||||
>
|
||||
<Box
|
||||
$flex="1"
|
||||
|
||||
@@ -40,7 +40,7 @@ export const ToolInvocationItem: React.FC<ToolInvocationItemProps> = ({
|
||||
$color="var(--c--theme--colors--greyscale-500)"
|
||||
$padding={{ all: 'sm' }}
|
||||
$radius="8px"
|
||||
$css="font-size: 0.9em;"
|
||||
$css="font-size: 0.9em; width: 100%; white-space: pre-wrap; word-wrap: break-word;"
|
||||
>
|
||||
{toolInvocation.state === 'result' ? (
|
||||
<Text>{`Parsing done: ${documentIdentifiers.join(', ')}`}</Text>
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import { useState } from 'react';
|
||||
import { useCallback, useRef, useState } from 'react';
|
||||
|
||||
interface SourceMetadata {
|
||||
title: string | null;
|
||||
@@ -9,109 +9,124 @@ interface SourceMetadata {
|
||||
|
||||
// Cache global pour éviter de refetch les mêmes URLs
|
||||
const metadataCache = new Map<string, SourceMetadata>();
|
||||
const fetchingUrls = new Set<string>();
|
||||
|
||||
export const useSourceMetadataCache = () => {
|
||||
const [cache, setCache] =
|
||||
useState<Map<string, SourceMetadata>>(metadataCache);
|
||||
const [, forceUpdate] = useState({});
|
||||
const updateCountRef = useRef(0);
|
||||
|
||||
const prefetchMetadata = async (url: string) => {
|
||||
// Si déjà en cache, ne rien faire
|
||||
if (metadataCache.has(url)) {
|
||||
return;
|
||||
const triggerUpdate = useCallback(() => {
|
||||
updateCountRef.current++;
|
||||
if (updateCountRef.current % 5 === 0) {
|
||||
forceUpdate({});
|
||||
}
|
||||
}, []);
|
||||
|
||||
// Marquer comme en cours de chargement
|
||||
metadataCache.set(url, {
|
||||
title: null,
|
||||
favicon: null,
|
||||
loading: true,
|
||||
error: false,
|
||||
});
|
||||
setCache(new Map(metadataCache));
|
||||
|
||||
try {
|
||||
if (!url.startsWith('http')) {
|
||||
metadataCache.set(url, {
|
||||
title: url,
|
||||
favicon: '📄',
|
||||
loading: false,
|
||||
error: false,
|
||||
});
|
||||
setCache(new Map(metadataCache));
|
||||
const prefetchMetadata = useCallback(
|
||||
async (url: string) => {
|
||||
if (metadataCache.has(url) || fetchingUrls.has(url)) {
|
||||
return;
|
||||
}
|
||||
|
||||
const parser = new DOMParser();
|
||||
fetchingUrls.add(url);
|
||||
|
||||
metadataCache.set(url, {
|
||||
title: null,
|
||||
favicon: null,
|
||||
loading: true,
|
||||
error: false,
|
||||
});
|
||||
triggerUpdate();
|
||||
|
||||
let response;
|
||||
try {
|
||||
response = await fetch(url, {
|
||||
mode: 'cors',
|
||||
headers: {
|
||||
'User-Agent': 'Mozilla/5.0 (compatible; ChatBot/1.0)',
|
||||
},
|
||||
if (!url.startsWith('http')) {
|
||||
metadataCache.set(url, {
|
||||
title: url,
|
||||
favicon: '📄',
|
||||
loading: false,
|
||||
error: false,
|
||||
});
|
||||
fetchingUrls.delete(url);
|
||||
triggerUpdate();
|
||||
return;
|
||||
}
|
||||
|
||||
const parser = new DOMParser();
|
||||
|
||||
let response;
|
||||
try {
|
||||
response = await fetch(url, {
|
||||
mode: 'cors',
|
||||
headers: {
|
||||
'User-Agent': 'Mozilla/5.0 (compatible; ChatBot/1.0)',
|
||||
},
|
||||
});
|
||||
} catch {
|
||||
// Si CORS échoue, utiliser juste le hostname
|
||||
metadataCache.set(url, {
|
||||
title: new URL(url).hostname,
|
||||
favicon: null,
|
||||
loading: false,
|
||||
error: false,
|
||||
});
|
||||
fetchingUrls.delete(url);
|
||||
triggerUpdate();
|
||||
return;
|
||||
}
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}`);
|
||||
}
|
||||
|
||||
const html = await response.text();
|
||||
const doc = parser.parseFromString(html, 'text/html');
|
||||
|
||||
// Récupérer le titre
|
||||
const pageTitle =
|
||||
doc.querySelector('title')?.textContent || new URL(url).hostname;
|
||||
|
||||
// Récupérer le favicon
|
||||
let faviconUrl =
|
||||
doc.querySelector('link[rel="icon"]')?.getAttribute('href') ||
|
||||
doc.querySelector('link[rel="shortcut icon"]')?.getAttribute('href');
|
||||
|
||||
if (!faviconUrl) {
|
||||
const urlObj = new URL(url);
|
||||
faviconUrl = `${urlObj.origin}/favicon.ico`;
|
||||
}
|
||||
|
||||
// Convertir les URLs relatives en absolues
|
||||
if (faviconUrl && !faviconUrl.startsWith('http')) {
|
||||
const urlObj = new URL(url);
|
||||
faviconUrl = new URL(faviconUrl, urlObj.origin).href;
|
||||
}
|
||||
|
||||
metadataCache.set(url, {
|
||||
title: pageTitle,
|
||||
favicon: faviconUrl || null,
|
||||
loading: false,
|
||||
error: false,
|
||||
});
|
||||
} catch {
|
||||
// Si CORS échoue, utiliser juste le hostname
|
||||
fetchingUrls.delete(url);
|
||||
triggerUpdate();
|
||||
} catch (err) {
|
||||
console.log('Error fetching metadata for:', url, err);
|
||||
metadataCache.set(url, {
|
||||
title: new URL(url).hostname,
|
||||
favicon: null,
|
||||
loading: false,
|
||||
error: false,
|
||||
error: true,
|
||||
});
|
||||
setCache(new Map(metadataCache));
|
||||
return;
|
||||
fetchingUrls.delete(url);
|
||||
triggerUpdate();
|
||||
}
|
||||
},
|
||||
[triggerUpdate],
|
||||
);
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP ${response.status}`);
|
||||
}
|
||||
const getMetadata = useCallback((url: string): SourceMetadata | undefined => {
|
||||
return metadataCache.get(url);
|
||||
}, []);
|
||||
|
||||
const html = await response.text();
|
||||
const doc = parser.parseFromString(html, 'text/html');
|
||||
|
||||
// Récupérer le titre
|
||||
const pageTitle =
|
||||
doc.querySelector('title')?.textContent || new URL(url).hostname;
|
||||
|
||||
// Récupérer le favicon
|
||||
let faviconUrl =
|
||||
doc.querySelector('link[rel="icon"]')?.getAttribute('href') ||
|
||||
doc.querySelector('link[rel="shortcut icon"]')?.getAttribute('href');
|
||||
|
||||
if (!faviconUrl) {
|
||||
const urlObj = new URL(url);
|
||||
faviconUrl = `${urlObj.origin}/favicon.ico`;
|
||||
}
|
||||
|
||||
// Convertir les URLs relatives en absolues
|
||||
if (faviconUrl && !faviconUrl.startsWith('http')) {
|
||||
const urlObj = new URL(url);
|
||||
faviconUrl = new URL(faviconUrl, urlObj.origin).href;
|
||||
}
|
||||
|
||||
metadataCache.set(url, {
|
||||
title: pageTitle,
|
||||
favicon: faviconUrl || null,
|
||||
loading: false,
|
||||
error: false,
|
||||
});
|
||||
setCache(new Map(metadataCache));
|
||||
} catch (err) {
|
||||
console.log('Error fetching metadata for:', url, err);
|
||||
metadataCache.set(url, {
|
||||
title: new URL(url).hostname,
|
||||
favicon: null,
|
||||
loading: false,
|
||||
error: true,
|
||||
});
|
||||
setCache(new Map(metadataCache));
|
||||
}
|
||||
};
|
||||
|
||||
const getMetadata = (url: string): SourceMetadata | undefined => {
|
||||
return cache.get(url);
|
||||
};
|
||||
|
||||
return { prefetchMetadata, getMetadata, cache };
|
||||
return { prefetchMetadata, getMetadata };
|
||||
};
|
||||
|
||||
@@ -54,7 +54,7 @@ export const Feedback = (_props: { buttonProps?: Partial<ButtonProps> }) => {
|
||||
/>
|
||||
|
||||
<FeedbackButton
|
||||
href="https://www.tchap.gouv.fr/#/room/!eAHyPLdVHMxNhKAbaC:agent.dinum.tchap.gouv.fr"
|
||||
href="https://tchap.gouv.fr/#/room/!eAHyPLdVHMxNhKAbaC:agent.dinum.tchap.gouv.fr?via=agent.dinum.tchap.gouv.fr&via=agent.culture.tchap.gouv.fr&via=agent.education.tchap.gouv.fr"
|
||||
icon={<TchapIcon />}
|
||||
title={t('Write on Tchap')}
|
||||
description={t('Direct exchange with our team')}
|
||||
|
||||
+4
-2
@@ -6,13 +6,15 @@ import { useChatPreferencesStore } from '@/features/chat/stores/useChatPreferenc
|
||||
|
||||
export const ButtonToggleLeftPanel = () => {
|
||||
const { t } = useTranslation();
|
||||
const { togglePanel } = useChatPreferencesStore();
|
||||
const { isPanelOpen, togglePanel } = useChatPreferencesStore();
|
||||
|
||||
return (
|
||||
<Button
|
||||
size="medium"
|
||||
onClick={() => togglePanel()}
|
||||
aria-label={t('Open the header menu')}
|
||||
aria-label={
|
||||
isPanelOpen ? t('Close the left panel') : t('Open the left panel')
|
||||
}
|
||||
color="primary-text"
|
||||
icon={<LeftPanelIcon />}
|
||||
className="--docs--button-toggle-panel"
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user