mirror of https://github.com/buster-so/buster.git
Merge branch 'evals' of https://github.com/buster-so/buster into evals
This commit is contained in:
commit
2ecc49eaaf
|
@ -0,0 +1,104 @@
|
|||
# Guidelines for Claude Code
|
||||
|
||||
## Code Quality Standards
|
||||
- Always clean up warnings shown by `cargo check` or similar commands
|
||||
- Remove unused dependencies from Cargo.toml files
|
||||
- Follow Rust best practices and idiomatic patterns
|
||||
- Ensure code passes linting and typechecking before completion
|
||||
- Use `cargo clippy` to identify and fix code quality issues
|
||||
- All code must pass `cargo check` before being merged into any other branch
|
||||
- Never merge code that has failing tests or compiler errors
|
||||
|
||||
## Project Structure
|
||||
- Maintain the workspace-based organization with multiple crates
|
||||
- Keep the API modular with separate libraries for specific functionality
|
||||
- Follow the established routing and module organization patterns
|
||||
- Design for modularity to enable unit testing and component isolation
|
||||
- Structure code to minimize dependencies and enable independent testing
|
||||
|
||||
## Error Handling
|
||||
- Use thiserror for defining custom error types
|
||||
- Implement From trait for error conversions
|
||||
- Use the ? operator for error propagation
|
||||
- Include descriptive error messages
|
||||
|
||||
## Code Style
|
||||
- Use 4-space indentation
|
||||
- Follow consistent naming conventions
|
||||
- Organize imports logically
|
||||
- Use modern Rust idioms
|
||||
- Write modular code with clear separation of concerns
|
||||
- Design components to be testable from the start
|
||||
|
||||
## Testing
|
||||
- Write unit tests alongside implementation code
|
||||
- Create integration tests for critical functionality
|
||||
- Use test fixtures for common test setup
|
||||
- All tests must pass before merging to other branches
|
||||
- Run `cargo test` locally to verify changes don't break existing functionality
|
||||
- Scope tests to the parts of the code being worked on during development (`cargo test -p <package_name>`)
|
||||
- Run the full test suite before submitting changes for review
|
||||
- Aim for high test coverage on critical components
|
||||
- Write tests that are meaningful and verify actual business logic
|
||||
|
||||
## Documentation
|
||||
- Document public interfaces and modules
|
||||
- Include clear comments for complex implementations
|
||||
- Keep documentation up to date with code changes
|
||||
|
||||
## Library Architecture
|
||||
|
||||
The Buster API is organized into several modular libraries, each with a specific purpose. When working with these libraries, maintain their architectural boundaries and respect their intended uses.
|
||||
|
||||
### Core Libraries
|
||||
|
||||
| Library | Purpose | Key Functionality |
|
||||
|---------|---------|-------------------|
|
||||
| `agents` | Provides agent functionality for interacting with LLMs | Agent interfaces, tool execution framework, specialized agents for data-related tasks |
|
||||
| `database` | Core database interaction and models | Connection pooling, models, schema definitions, database helpers |
|
||||
| `handlers` | Request handlers for the REST API | Endpoint implementations, request validation, response formatting |
|
||||
| `query_engine` | Data source connection and query execution | Database connections, query execution, data transformation |
|
||||
| `streaming` | Streaming parser for LLM responses | JSON stream parsing, partial data processing, processor registry |
|
||||
|
||||
### Supporting Libraries
|
||||
|
||||
| Library | Purpose | Key Functionality |
|
||||
|---------|---------|-------------------|
|
||||
| `braintrust` | Client for Braintrust API integration | Spans and traces for monitoring AI application performance |
|
||||
| `litellm` | Client for LLM providers | Unified interface to various LLM providers, message handling |
|
||||
| `middleware` | HTTP middleware components | Authentication, CORS, request validation |
|
||||
| `sharing` | Asset sharing and permission management | Permission checks, asset sharing, user lookup |
|
||||
|
||||
### Library Dependency Map
|
||||
|
||||
- `handlers` → depends on most other libraries to implement API endpoints
|
||||
- `agents` → depends on `litellm`, `database`, `query_engine`
|
||||
- `query_engine` → depends on `database`
|
||||
- `sharing` → depends on `database`
|
||||
- `middleware` → depends on `database`
|
||||
|
||||
When working on a feature, start by identifying which library or libraries are responsible for the functionality, and respect the existing architecture patterns.
|
||||
|
||||
### Library Management Guidelines
|
||||
|
||||
- When adding a new library to the codebase, always:
|
||||
- Update this CLAUDE.md file with:
|
||||
- The library's name, purpose, and key functionality
|
||||
- Its placement in the appropriate category (Core or Supporting)
|
||||
- Any dependencies it has on other libraries
|
||||
- Update the dependency map if needed
|
||||
- Create a CLAUDE.md in the library's root directory that includes:
|
||||
- Detailed explanation of the library's purpose and role
|
||||
- Overview of its internal organization and key modules
|
||||
- Usage patterns and implementation guidance
|
||||
- Dependencies and relationships with other libraries
|
||||
- Code navigation tips for the agent
|
||||
- Testing guidelines specific to this library
|
||||
|
||||
- When significantly changing a library's purpose or functionality:
|
||||
- Update its description in this document
|
||||
- Update the library's CLAUDE.md with the new information
|
||||
- Update any dependent libraries' documentation if the changes affect their usage
|
||||
|
||||
- Maintain architectural boundaries between libraries to prevent circular dependencies
|
||||
- Each library should have clear and well-documented public interfaces
|
|
@ -5,6 +5,168 @@ alwaysApply: false
|
|||
---
|
||||
# Testing Rules and Best Practices
|
||||
|
||||
## Testing Organization Best Practices
|
||||
|
||||
### Recommended Mocking Libraries
|
||||
|
||||
For consistent mocking across the codebase, use these libraries:
|
||||
|
||||
| Library | Purpose | Use Case |
|
||||
|---------|---------|----------|
|
||||
| **mockito** | HTTP service mocking | Mocking external API calls |
|
||||
| **mockall** | Trait/struct mocking | Mocking database and service interfaces |
|
||||
| **mock-it** | Simple mocking | Quick mocks for simple interfaces |
|
||||
| **wiremock** | Advanced HTTP mocking | Complex API scenarios with extended matching |
|
||||
|
||||
Always prefer dependency injection patterns to enable easy mocking of dependencies.
|
||||
|
||||
### Library Code Testing Structure
|
||||
- **Unit Tests**: Include unit tests inside library source files using `#[cfg(test)]` modules
|
||||
- **Integration Tests**: Place integration tests in the lib's `tests/` directory
|
||||
- **Test Utils**: Create a private `test_utils.rs` module in each lib for test utilities specific to that lib
|
||||
- **Scope Isolation**: Test utilities should only be shared within their scope (e.g., libs/database, libs/agents)
|
||||
- **Conditional Testing**: Use conditional compilation or runtime checks to make tests non-blocking (e.g., skip tests if a required API key is missing)
|
||||
|
||||
### API (src) Testing Structure
|
||||
- Maintain API-specific tests in the `/api/tests/` directory
|
||||
- Organize integration tests to mirror the API's route structure
|
||||
- Keep API-specific test utilities in the shared test common directory
|
||||
|
||||
### Testing Configuration
|
||||
- Use cargo features to make running specific test groups easier
|
||||
- Consider using cargo-nextest for parallel test execution
|
||||
- Add test tags to group tests by functionality or component
|
||||
|
||||
### Running Tests in Different Scopes
|
||||
|
||||
- **Run all tests**: `cargo test`
|
||||
- **Run library-specific tests**: `cargo test -p <library-name>`
|
||||
- **Run a specific test**: `cargo test <test_name>`
|
||||
- **Run tests with specific pattern**: `cargo test -- <pattern>`
|
||||
- **Run tests conditionally**: Use features to enable/disable test categories
|
||||
|
||||
```toml
|
||||
# In Cargo.toml
|
||||
[features]
|
||||
integration-tests = []
|
||||
unit-tests = []
|
||||
```
|
||||
|
||||
```rust
|
||||
// In code
|
||||
#[cfg(test)]
|
||||
#[cfg(feature = "integration-tests")]
|
||||
mod integration_tests {
|
||||
// Test code...
|
||||
}
|
||||
```
|
||||
|
||||
Then, run: `cargo test --features integration-tests`
|
||||
|
||||
### Examples of Proper Test Organization
|
||||
|
||||
#### Example: Library Unit Tests (`#[cfg(test)]` Module)
|
||||
```rust
|
||||
// File: libs/braintrust/src/helpers.rs
|
||||
|
||||
// Implementation code...
|
||||
|
||||
/// Get system message from a stored prompt
|
||||
pub async fn get_prompt_system_message(client: &BraintrustClient, prompt_id: &str) -> Result<String> {
|
||||
// Function implementation...
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::env;
|
||||
use dotenv::dotenv;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_prompt_system_message() -> Result<()> {
|
||||
// Load environment variables
|
||||
dotenv().ok();
|
||||
|
||||
// Skip test if no API key is available (non-blocking)
|
||||
if env::var("BRAINTRUST_API_KEY").is_err() {
|
||||
println!("Skipping test_get_prompt_system_message: No API key available");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Test implementation...
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Example: Library Integration Tests (Separate Directory)
|
||||
```rust
|
||||
// File: libs/handlers/tests/metrics/delete_metric_test.rs
|
||||
|
||||
use handlers::metrics::delete_metric_handler;
|
||||
use database::models::MetricFile;
|
||||
use crate::common::fixtures::create_test_metric;
|
||||
use anyhow::Result;
|
||||
|
||||
/// Integration test for the delete_metric_handler
|
||||
#[tokio::test]
|
||||
async fn test_delete_metric_integration() -> Result<()> {
|
||||
// Setup test environment
|
||||
let test_db = setup_test_db().await?;
|
||||
|
||||
// Create and insert test data
|
||||
let test_metric = create_test_metric(&test_db).await?;
|
||||
|
||||
// Test the handler functionality
|
||||
let result = delete_metric_handler(&test_metric.id, &test_db.user_id).await?;
|
||||
|
||||
// Verify results
|
||||
assert!(result.deleted_at.is_some());
|
||||
|
||||
// Clean up test data
|
||||
test_db.cleanup().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
#### Example: Library-Specific Test Utilities
|
||||
```rust
|
||||
// File: libs/database/src/test_utils.rs
|
||||
|
||||
use crate::models::{User, Organization};
|
||||
use crate::pool::PgPool;
|
||||
use uuid::Uuid;
|
||||
use anyhow::Result;
|
||||
|
||||
/// Test utilities specific to the database library
|
||||
pub struct DatabaseTestUtils {
|
||||
pub pool: PgPool,
|
||||
pub test_id: String,
|
||||
}
|
||||
|
||||
impl DatabaseTestUtils {
|
||||
pub async fn new() -> Result<Self> {
|
||||
// Initialize test database connection
|
||||
let pool = create_test_pool().await?;
|
||||
let test_id = Uuid::new_v4().to_string();
|
||||
|
||||
Ok(Self { pool, test_id })
|
||||
}
|
||||
|
||||
pub async fn create_test_user(&self) -> Result<User> {
|
||||
// Create a test user in the database
|
||||
// ...
|
||||
}
|
||||
|
||||
pub async fn cleanup(&self) -> Result<()> {
|
||||
// Clean up test data
|
||||
// ...
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## General Testing Guidelines
|
||||
- All tests must be async and use tokio test framework
|
||||
- Tests should be well-documented with clear test case descriptions and expected outputs
|
||||
|
@ -21,7 +183,11 @@ alwaysApply: false
|
|||
## Unit Tests
|
||||
- Unit tests should be inline with the code they are testing using `#[cfg(test)]` modules
|
||||
- Each public function should have corresponding unit tests
|
||||
- Mock external dependencies using mockito for HTTP calls
|
||||
- **Important**: Unit tests should NEVER connect to external services or databases
|
||||
- Mock all external dependencies:
|
||||
- Use mockito for HTTP services
|
||||
- Use trait-based mocks for database operations
|
||||
- Create mock implementations of dependencies
|
||||
- **IMPORTANT**: Always use the async version of mockito:
|
||||
```rust
|
||||
// Correct async approach
|
||||
|
@ -36,19 +202,143 @@ alwaysApply: false
|
|||
- Test both success and error cases
|
||||
- Test edge cases and boundary conditions
|
||||
- Structure unit tests to maximize code coverage
|
||||
- Use dependency injection to make code easily testable with mocks
|
||||
|
||||
## Integration Tests
|
||||
- Integration tests for the main application should be placed in the `/tests` directory
|
||||
- **For libraries in the `libs/` directory**:
|
||||
- Integration tests should be within each library in their own `tests/` directory
|
||||
- Each library should be responsible for testing its own public API
|
||||
- Follow the Rust convention of placing integration tests in the library's `tests/` directory
|
||||
- Organize integration tests to mirror the main codebase structure
|
||||
- Each major feature/resource should have its own test file
|
||||
- Test the interaction between multiple components
|
||||
- Use real dependencies when possible, mock only what's necessary
|
||||
- Include end-to-end workflow tests
|
||||
- Test realistic user workflows rather than individual functions
|
||||
|
||||
Integration tests are specifically designed to verify the interaction with external services and should be separate from the main codebase.
|
||||
|
||||
### Integration Test Structure
|
||||
|
||||
- **Location**: Integration tests should always be in a separate `tests/` directory, never mixed with the main code
|
||||
- **External Interaction**: Unlike unit tests, integration tests are explicitly designed to interact with external services (databases, APIs, etc.)
|
||||
- **Configuration**:
|
||||
- Configuration should come from environment variables via `.env` files
|
||||
- Use `dotenv` to load environment variables during test setup
|
||||
- Example: Database connection parameters should come from `.env`
|
||||
- Prefer `.env.test` for test-specific configurations
|
||||
|
||||
```rust
|
||||
// Example of proper integration test setup
|
||||
use dotenv::dotenv;
|
||||
use diesel_async::AsyncPgConnection;
|
||||
use diesel_async::pooled_connection::AsyncDieselConnectionManager;
|
||||
use deadpool_diesel::postgres::Pool;
|
||||
|
||||
// Setup function that loads from environment
|
||||
async fn setup_test_environment() -> Result<TestContext> {
|
||||
// Load environment variables, preferring .env.test if available
|
||||
if std::path::Path::new(".env.test").exists() {
|
||||
dotenv::from_filename(".env.test").ok();
|
||||
} else {
|
||||
dotenv().ok();
|
||||
}
|
||||
|
||||
// Create database pool from environment variables
|
||||
let database_url = std::env::var("DATABASE_URL")
|
||||
.expect("DATABASE_URL must be set for integration tests");
|
||||
|
||||
let config = AsyncDieselConnectionManager::<AsyncPgConnection>::new(database_url);
|
||||
let pool = Pool::builder(config).max_size(5).build()?;
|
||||
|
||||
// Return test context with real connections to external services
|
||||
Ok(TestContext { pool, ... })
|
||||
}
|
||||
```
|
||||
|
||||
### Library-Specific Integration Tests
|
||||
|
||||
- Each library should have its own integration tests in its `tests/` directory
|
||||
- These tests should focus ONLY on the library's public API
|
||||
- Integration tests should NOT cross library boundaries
|
||||
- Example structure:
|
||||
|
||||
```
|
||||
libs/
|
||||
├── database/
|
||||
│ ├── src/
|
||||
│ └── tests/ # Database-specific integration tests
|
||||
│ ├── common/ # Database test utilities
|
||||
│ └── models/ # Tests for database models
|
||||
├── handlers/
|
||||
│ ├── src/
|
||||
│ └── tests/ # Handler-specific integration tests
|
||||
│ ├── common/ # Handler test utilities
|
||||
│ └── users/ # Tests for user handlers
|
||||
```
|
||||
|
||||
### API Integration Tests
|
||||
|
||||
- API tests should be in the main `/api/tests/` directory
|
||||
- Structure should mirror the API routes
|
||||
- Focus on end-to-end testing of the entire API
|
||||
- Example structure:
|
||||
|
||||
```
|
||||
api/
|
||||
├── src/
|
||||
└── tests/
|
||||
├── common/ # Shared API test utilities
|
||||
└── integration/ # Organized by API resource
|
||||
├── users/ # User endpoint tests
|
||||
├── threads/ # Thread endpoint tests
|
||||
└── messages/ # Message endpoint tests
|
||||
```
|
||||
|
||||
### Modular Test Utilities
|
||||
|
||||
- **Scope Isolation**: Test utilities should be scoped to their specific domain:
|
||||
- Database utilities in `libs/database/tests/common/`
|
||||
- Handler utilities in `libs/handlers/tests/common/`
|
||||
- API utilities in `/api/tests/common/`
|
||||
|
||||
- **Modular Structure**: Organize test utilities by function:
|
||||
|
||||
```
|
||||
tests/common/
|
||||
├── mod.rs # Re-exports common utilities
|
||||
├── db.rs # Database test helpers
|
||||
├── http.rs # HTTP client testing
|
||||
├── auth.rs # Authentication test helpers
|
||||
└── fixtures/ # Test data fixtures
|
||||
├── mod.rs # Re-exports all fixtures
|
||||
├── users.rs # User test data
|
||||
└── messages.rs # Message test data
|
||||
```
|
||||
|
||||
- **No Global Utilities**: Avoid creating global test utilities shared across all components
|
||||
- **Library-Specific Fixtures**: Test fixtures should be specific to their domain
|
||||
|
||||
### Test Data Isolation
|
||||
|
||||
- Each test should create its own isolated data
|
||||
- Use unique identifiers (UUIDs) to mark test data
|
||||
- Clean up after tests complete
|
||||
- Use test_id pattern to track and clean up test data:
|
||||
|
||||
```rust
|
||||
#[tokio::test]
|
||||
async fn test_user_creation() -> Result<()> {
|
||||
let test_id = Uuid::new_v4().to_string();
|
||||
let pool = get_db_pool().await?;
|
||||
|
||||
// Create test data with test_id marker
|
||||
let test_user = UserBuilder::new()
|
||||
.with_email("test@example.com")
|
||||
.with_test_id(&test_id) // Mark this as test data
|
||||
.build();
|
||||
|
||||
// Insert test data
|
||||
test_user.insert(&pool).await?;
|
||||
|
||||
// Run test assertions...
|
||||
|
||||
// Clean up all data with this test_id
|
||||
cleanup_test_data(&pool, &test_id).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Integration Test Setup Best Practices
|
||||
|
||||
|
@ -244,11 +534,362 @@ mod tests {
|
|||
}
|
||||
```
|
||||
|
||||
## Mocking Guidelines
|
||||
- Use mockito for HTTP service mocks
|
||||
- Create mock responses that match real API responses
|
||||
- Include both successful and error responses in mocks
|
||||
- Clean up mocks after tests complete
|
||||
## Mocking Guidelines for Async Testing
|
||||
|
||||
### Mockito for HTTP Service Mocking
|
||||
|
||||
Mockito is the primary tool for mocking HTTP services in tests. When using mockito in an async context:
|
||||
|
||||
```rust
|
||||
// CORRECT: Create async mockito server
|
||||
#[tokio::test]
|
||||
async fn test_http_client() -> Result<()> {
|
||||
// Always use the async version for tokio compatibility
|
||||
let server = mockito::Server::new_async().await;
|
||||
|
||||
// Setup the mock with expected request and response
|
||||
let mock = server
|
||||
.mock("GET", "/api/data")
|
||||
.with_status(200)
|
||||
.with_header("content-type", "application/json")
|
||||
.with_body(r#"{"key": "value"}"#)
|
||||
.create();
|
||||
|
||||
// Use the server's URL with your client
|
||||
let client = YourHttpClient::new(&server.url());
|
||||
let response = client.get_data().await?;
|
||||
|
||||
// Verify the mock was called as expected
|
||||
mock.assert();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
#### Mockito Best Practices:
|
||||
|
||||
- **Always use `Server::new_async()`** instead of `Server::new()` for tokio compatibility
|
||||
- **Match complex requests** with matchers:
|
||||
```rust
|
||||
// Match JSON request bodies
|
||||
server.mock("POST", "/api/users")
|
||||
.match_body(mockito::Matcher::Json(json!({"name": "Test User"})))
|
||||
.with_status(201)
|
||||
.create();
|
||||
|
||||
// Match headers
|
||||
server.mock("GET", "/api/protected")
|
||||
.match_header("authorization", "Bearer token")
|
||||
.with_status(200)
|
||||
.create();
|
||||
```
|
||||
- **Test different response scenarios** (success, errors, timeouts):
|
||||
```rust
|
||||
// Success response
|
||||
let success_mock = server.mock("GET", "/api/resource/1")
|
||||
.with_status(200)
|
||||
.with_body(r#"{"id": "1", "name": "Resource"}"#)
|
||||
.create();
|
||||
|
||||
// Error response
|
||||
let error_mock = server.mock("GET", "/api/resource/999")
|
||||
.with_status(404)
|
||||
.with_body(r#"{"error": "Not found"}"#)
|
||||
.create();
|
||||
|
||||
// Timeout simulation
|
||||
let timeout_mock = server.mock("GET", "/api/slow")
|
||||
.with_delay(std::time::Duration::from_secs(10))
|
||||
.with_status(200)
|
||||
.create();
|
||||
```
|
||||
- **Create reusable mock setup functions** for common patterns
|
||||
|
||||
### Additional Mocking Libraries for Async Rust
|
||||
|
||||
#### 1. Mock_it for trait mocking
|
||||
|
||||
For mocking traits and interfaces:
|
||||
|
||||
```rust
|
||||
use mock_it::Mock;
|
||||
|
||||
#[async_trait]
|
||||
pub trait Database {
|
||||
async fn get_user(&self, id: &str) -> Result<User>;
|
||||
}
|
||||
|
||||
// Create a mock implementation
|
||||
let mut db_mock = Mock::new();
|
||||
|
||||
// Set up expectations with async behavior
|
||||
db_mock.expect_call(matching!("get_user", arg if arg == "123"))
|
||||
.returns_async(Ok(User {
|
||||
id: "123".to_string(),
|
||||
name: "Test User".to_string(),
|
||||
}));
|
||||
|
||||
// Test using the mock
|
||||
let result = db_mock.get_user("123").await?;
|
||||
assert_eq!(result.name, "Test User");
|
||||
```
|
||||
|
||||
#### 2. Async-std's Async-Mock
|
||||
|
||||
For async-std runtime testing:
|
||||
|
||||
```rust
|
||||
use async_std::test;
|
||||
use async_mock::AsyncMock;
|
||||
|
||||
#[derive(AsyncMock)]
|
||||
#[async_mock(UserRepository)]
|
||||
pub trait UserRepository {
|
||||
async fn find_user(&self, id: &str) -> Result<User>;
|
||||
}
|
||||
|
||||
#[test]
|
||||
async fn test_user_service() {
|
||||
// Create a mock repository
|
||||
let mock_repo = MockUserRepository::new();
|
||||
|
||||
// Setup expectations
|
||||
mock_repo.expect_find_user()
|
||||
.with(eq("123"))
|
||||
.returns(Ok(User { id: "123".to_string(), name: "User" }));
|
||||
|
||||
// Test service with mock
|
||||
let service = UserService::new(mock_repo);
|
||||
let user = service.get_user_data("123").await?;
|
||||
|
||||
assert_eq!(user.name, "User");
|
||||
}
|
||||
```
|
||||
|
||||
#### 3. WireMock for Complex HTTP Mocking
|
||||
|
||||
For extensive API mocking scenarios:
|
||||
|
||||
```rust
|
||||
use wiremock::{MockServer, Mock, ResponseTemplate};
|
||||
use wiremock::matchers::{method, path};
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_api_client() -> Result<()> {
|
||||
// Start mock server
|
||||
let mock_server = MockServer::start().await;
|
||||
|
||||
// Setup mock
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/users"))
|
||||
.respond_with(ResponseTemplate::new(200)
|
||||
.set_body_json(json!([
|
||||
{"id": "1", "name": "User 1"},
|
||||
{"id": "2", "name": "User 2"}
|
||||
]))
|
||||
)
|
||||
.mount(&mock_server)
|
||||
.await;
|
||||
|
||||
// Create client using mock URL
|
||||
let client = ApiClient::new(&mock_server.uri());
|
||||
|
||||
// Test the client
|
||||
let users = client.list_users().await?;
|
||||
assert_eq!(users.len(), 2);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
#### 4. Tower's `mock` Module for Service Mocking
|
||||
|
||||
For mocking Tower services:
|
||||
|
||||
```rust
|
||||
use tower_test::{mock, assert_request_eq};
|
||||
use tower::Service;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tower_service() {
|
||||
// Create mock service
|
||||
let (mut mock, handle) = mock::pair::<Request<()>, Response<()>>();
|
||||
|
||||
// Spawn a task that processes requests
|
||||
tokio::spawn(async move {
|
||||
// Send successful response
|
||||
if let Ok(request) = handle.recv().await {
|
||||
handle.send_response(
|
||||
Response::builder()
|
||||
.status(200)
|
||||
.body(())
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// Create your service that uses the mock
|
||||
let service = YourService::new(mock);
|
||||
|
||||
// Make request through your service
|
||||
let response = service
|
||||
.call(Request::get("/test").body(()).unwrap())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), 200);
|
||||
}
|
||||
```
|
||||
|
||||
### Mocking Database Access for Unit Tests
|
||||
|
||||
For unit tests, completely mock the database access rather than using real connections:
|
||||
|
||||
```rust
|
||||
use mock_it::Mock;
|
||||
use async_trait::async_trait;
|
||||
use uuid::Uuid;
|
||||
|
||||
// Define a repository trait that can be mocked
|
||||
#[async_trait]
|
||||
pub trait UserRepository {
|
||||
async fn find_user_by_id(&self, id: &str) -> Result<Option<User>>;
|
||||
async fn create_user(&self, email: &str, name: &str) -> Result<User>;
|
||||
async fn delete_user(&self, id: &str) -> Result<()>;
|
||||
}
|
||||
|
||||
// In unit tests, create a mock implementation
|
||||
#[tokio::test]
|
||||
async fn test_user_service_with_mocked_db() -> Result<()> {
|
||||
// Create mock repository
|
||||
let mut repo_mock = Mock::new();
|
||||
|
||||
// Setup expectations for database operations
|
||||
let test_user = User {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
email: "test@example.com".to_string(),
|
||||
name: "Test User".to_string(),
|
||||
created_at: chrono::Utc::now(),
|
||||
};
|
||||
|
||||
// Mock the find_user_by_id method
|
||||
repo_mock.expect_call(matching!("find_user_by_id", arg if arg == "123"))
|
||||
.returns_async(Ok(Some(test_user.clone())));
|
||||
|
||||
// Create the service with the mocked repository
|
||||
let service = UserService::new(repo_mock);
|
||||
|
||||
// Test the service logic
|
||||
let user = service.get_user_profile("123").await?;
|
||||
|
||||
// Assertions to verify service logic
|
||||
assert_eq!(user.email, "test@example.com");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Mocking Diesel Models for Unit Tests
|
||||
|
||||
For mocking Diesel model operations:
|
||||
|
||||
```rust
|
||||
use mockall::predicate::*;
|
||||
use mockall::mock;
|
||||
|
||||
// Create a mock for database operations
|
||||
mock! {
|
||||
pub DieselDb {
|
||||
async fn find_user(&self, id: &str) -> Result<Option<User>>;
|
||||
async fn create_user(&self, data: NewUser) -> Result<User>;
|
||||
}
|
||||
}
|
||||
|
||||
// User service that depends on database
|
||||
struct UserService {
|
||||
db: Box<dyn DieselDbTrait>,
|
||||
}
|
||||
|
||||
impl UserService {
|
||||
fn new(db: impl DieselDbTrait + 'static) -> Self {
|
||||
Self { db: Box::new(db) }
|
||||
}
|
||||
|
||||
async fn get_user(&self, id: &str) -> Result<Option<User>> {
|
||||
self.db.find_user(id).await
|
||||
}
|
||||
}
|
||||
|
||||
// Unit test with mocked database
|
||||
#[tokio::test]
|
||||
async fn test_get_user() -> Result<()> {
|
||||
// Create mock DB
|
||||
let mut mock_db = MockDieselDb::new();
|
||||
|
||||
// Setup expectations
|
||||
let test_user = User {
|
||||
id: "user123".to_string(),
|
||||
name: "Test User".to_string(),
|
||||
email: "test@example.com".to_string(),
|
||||
created_at: chrono::Utc::now(),
|
||||
};
|
||||
|
||||
mock_db.expect_find_user()
|
||||
.with(eq("user123"))
|
||||
.times(1)
|
||||
.returning(move |_| Ok(Some(test_user.clone())));
|
||||
|
||||
// Create service with mock
|
||||
let service = UserService::new(mock_db);
|
||||
|
||||
// Test service method
|
||||
let user = service.get_user("user123").await?;
|
||||
|
||||
// Verify results
|
||||
assert!(user.is_some());
|
||||
let user = user.unwrap();
|
||||
assert_eq!(user.email, "test@example.com");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Combining Different Mocking Approaches
|
||||
|
||||
For complex systems, combine different mocking libraries:
|
||||
|
||||
```rust
|
||||
#[tokio::test]
|
||||
async fn test_complex_system() -> Result<()> {
|
||||
// Set up HTTP mock
|
||||
let server = mockito::Server::new_async().await;
|
||||
|
||||
// Set up mock responses
|
||||
let api_mock = server
|
||||
.mock("GET", "/api/data")
|
||||
.with_status(200)
|
||||
.with_body(r#"{"data": "value"}"#)
|
||||
.create();
|
||||
|
||||
// Create mock database implementation
|
||||
let mut db_mock = Mock::new();
|
||||
db_mock.expect_call(matching!("save_data", _))
|
||||
.returns_async(Ok(()));
|
||||
|
||||
// Initialize system with mocks
|
||||
let client = HttpClient::new(&server.url());
|
||||
let system = YourSystem::new(client, db_mock);
|
||||
|
||||
// Test system behavior
|
||||
let result = system.process_and_save().await?;
|
||||
assert!(result.is_success());
|
||||
|
||||
// Verify HTTP mock was called
|
||||
api_mock.assert();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
## Error Testing
|
||||
- Test error conditions and error handling
|
||||
|
|
|
@ -0,0 +1,86 @@
|
|||
# Buster API Libraries
|
||||
|
||||
This directory contains the modular libraries that make up the Buster API. Each library has a specific purpose and follows the architectural principles outlined in the main CLAUDE.md file.
|
||||
|
||||
## Library CLAUDE.md Template
|
||||
|
||||
Below is a template that should be used when creating CLAUDE.md files for each library. Copy this template into each library's directory and fill it out with the specific details for that library.
|
||||
|
||||
```markdown
|
||||
# [Library Name] - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
Brief 1-2 sentence description of what this library does and its role in the overall system.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- List major capabilities provided by this library
|
||||
- Include key modules and their purposes
|
||||
- Highlight public interfaces and how they should be used
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── module1.rs - Purpose of this module
|
||||
├── module2/
|
||||
│ ├── submodule1.rs - Purpose of this submodule
|
||||
│ └── mod.rs
|
||||
├── types.rs - Core types used throughout the library
|
||||
└── lib.rs - Public exports and library documentation
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `module1`: Detailed explanation of what this module does
|
||||
- `module2`: Explanation of this module's functionality
|
||||
- `types`: Description of the core types and why they're designed this way
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
// Simple example of how to use this library
|
||||
use library_name::SomeType;
|
||||
|
||||
fn example() -> Result<(), Error> {
|
||||
let instance = SomeType::new()?;
|
||||
instance.do_something()?;
|
||||
// ...
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Describe typical usage patterns
|
||||
- Include best practices for using this library
|
||||
- Note any gotchas or non-obvious behavior
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- List library dependencies within the codebase and why they're needed
|
||||
|
||||
- **External Dependencies**:
|
||||
- List key external crates and their purposes in this library
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Important entry points to start exploring the code
|
||||
- Key type definitions to understand
|
||||
- Relationships between important components
|
||||
- How to trace execution flow through the library
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Guidance on how to test this library
|
||||
- Any specific test utilities or fixtures provided
|
||||
- How to run just the tests for this library: `cargo test -p library_name`
|
||||
```
|
||||
|
||||
## Existing Libraries
|
||||
|
||||
See the main [CLAUDE.md](../../CLAUDE.md) file for a complete list of libraries and their purposes.
|
|
@ -0,0 +1,112 @@
|
|||
# Agents Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Agents library provides a framework for creating and using AI agents that can interact with LLMs and execute tools. It serves as the foundation for implementing autonomous agents that can perform complex tasks like data analysis, dashboard creation, and planning.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Agent interfaces for standardized LLM interaction
|
||||
- Tool execution framework for defining and running agent tools
|
||||
- Specialized agents for data-related tasks
|
||||
- Tool categorization system for organizing agent capabilities
|
||||
- Utilities for agent message handling and response processing
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── agent.rs - Core agent trait definitions and implementations
|
||||
├── agents/
|
||||
│ ├── buster_super_agent.rs - Buster's primary agent implementation
|
||||
│ └── mod.rs
|
||||
├── models/
|
||||
│ ├── types.rs - Agent-related data structures
|
||||
│ └── mod.rs
|
||||
├── tools/
|
||||
│ ├── categories/ - Tool categorization
|
||||
│ │ ├── agents_as_tools/ - Using agents as tools
|
||||
│ │ ├── file_tools/ - File manipulation tools
|
||||
│ │ ├── planning_tools/ - Planning and strategy tools
|
||||
│ │ └── mod.rs
|
||||
│ ├── executor.rs - Tool execution machinery
|
||||
│ └── mod.rs
|
||||
├── utils/
|
||||
│ ├── tools.rs - Tool utility functions
|
||||
│ └── mod.rs
|
||||
└── lib.rs - Public exports and library documentation
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `agent`: Defines the core `Agent` trait and related functionality that all agents implement
|
||||
- `agents`: Contains concrete agent implementations, including the primary BusterSuperAgent
|
||||
- `models`: Data structures and types used throughout the agent system
|
||||
- `tools`: The tool execution framework, including the `ToolExecutor` trait
|
||||
- `utils`: Utility functions for agent operations and tool management
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use agents::{Agent, BusterSuperAgent};
|
||||
use litellm::AgentMessage;
|
||||
|
||||
async fn example_agent_use() -> Result<(), anyhow::Error> {
|
||||
// Create a new agent
|
||||
let agent = BusterSuperAgent::new(config).await?;
|
||||
|
||||
// Send a message to the agent
|
||||
let messages = vec![
|
||||
AgentMessage::user("Analyze the sales data for Q1")
|
||||
];
|
||||
|
||||
// Execute the agent and process the response
|
||||
let response = agent.execute(messages).await?;
|
||||
|
||||
// Process agent response...
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Agents are typically created once and reused for multiple interactions
|
||||
- Tools are registered with agents at creation time
|
||||
- Agent responses should be parsed for both content and tool call requests
|
||||
- Chain agent calls together for multi-step reasoning
|
||||
- Tool implementations should be focused, performing one specific task
|
||||
- Error handling flows up from tools to the agent executor
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- `litellm`: Used for LLM provider communication
|
||||
- `database`: Used for persisting agent state and accessing data
|
||||
- `query_engine`: Used by data-related tools to execute queries
|
||||
- `braintrust`: Used for agent performance monitoring
|
||||
|
||||
- **External Dependencies**:
|
||||
- `serde_json`: Used for parsing and generating JSON for agent communication
|
||||
- `anyhow` and `thiserror`: Used for error handling
|
||||
- `async-trait`: Used for async trait implementations
|
||||
- `tokio`: Used for async runtime
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see what's exported and main entry points
|
||||
- The `Agent` trait in `agent.rs` defines the core agent interface
|
||||
- `BusterSuperAgent` in `agents/buster_super_agent.rs` is the main implementation
|
||||
- Tool categories in `tools/categories/` are organized by functional area
|
||||
- The `ToolExecutor` trait in `tools/mod.rs` is implemented by various tool types
|
||||
- New tools should be added to the appropriate category in `tools/categories/`
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Unit tests focus on individual tool functionality
|
||||
- Mock the LLM responses using test fixtures
|
||||
- Use in-memory database for integration tests
|
||||
- Run tests with: `cargo test -p agents`
|
||||
- Test tools individually before testing complete agent flows
|
||||
- Create mock implementations of `Agent` trait for testing dependent code
|
|
@ -0,0 +1,94 @@
|
|||
# Braintrust Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Braintrust library provides a client for interacting with the Braintrust API, enabling the tracking and logging of AI application performance through spans and traces. It serves as the monitoring and observability layer for LLM interactions.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Braintrust API client for logging AI application performance
|
||||
- Span and trace creation for measuring LLM operations
|
||||
- Event and metrics tracking
|
||||
- Prompt and response logging
|
||||
- Structured metadata collection for AI operations
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── client.rs - BraintrustClient implementation
|
||||
├── types.rs - Core data structures for API interactions
|
||||
├── trace.rs - TraceBuilder for creating trace spans
|
||||
├── helpers.rs - Utility functions and helpers
|
||||
└── lib.rs - Public exports and documentation
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `client`: Implements the `BraintrustClient` for API communication
|
||||
- `types`: Defines data structures like `Span`, `Metrics`, and `EventPayload`
|
||||
- `trace`: Provides the `TraceBuilder` for creating and managing trace spans
|
||||
- `helpers`: Utility functions for working with the Braintrust API
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use braintrust::{BraintrustClient, TraceBuilder};
|
||||
|
||||
async fn example_usage() -> Result<(), anyhow::Error> {
|
||||
// Create a client
|
||||
let client = BraintrustClient::new("YOUR_API_KEY").await?;
|
||||
|
||||
// Create a trace
|
||||
let trace = TraceBuilder::new("example-trace")
|
||||
.with_metadata("user_id", "123")
|
||||
.build();
|
||||
|
||||
// Log a span with the client
|
||||
client.log_span(trace, "llm-call", |span| {
|
||||
// Your LLM call logic here
|
||||
span.with_input("What is the capital of France?")
|
||||
.with_output("Paris is the capital of France.")
|
||||
}).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Initialize the `BraintrustClient` once and reuse it throughout the application
|
||||
- Create traces for logical units of work (e.g., a chat session)
|
||||
- Use spans to measure individual operations within a trace (e.g., an LLM call)
|
||||
- Add metadata to traces and spans for better filtering and analysis
|
||||
- Properly handle errors from the Braintrust API
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- None - this is a standalone library that other libraries depend on
|
||||
|
||||
- **External Dependencies**:
|
||||
- `reqwest`: For making HTTP requests to the Braintrust API
|
||||
- `serde_json`: For JSON serialization/deserialization
|
||||
- `tokio`: For async runtime
|
||||
- `uuid`: For generating unique IDs
|
||||
- `chrono`: For timestamp handling
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see the exported API
|
||||
- The `BraintrustClient` in `client.rs` is the main entry point for usage
|
||||
- `types.rs` contains all the data structures used for API interactions
|
||||
- `trace.rs` contains the `TraceBuilder` for creating traces
|
||||
- When adding new metrics or events, extend the existing types in `types.rs`
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Mock the HTTP responses in unit tests using `mockito`
|
||||
- Test error handling by simulating API errors
|
||||
- Validate the correct structure of API requests in tests
|
||||
- Run tests with: `cargo test -p braintrust`
|
||||
- Use the examples in the `examples/` directory as reference implementations
|
|
@ -1,117 +0,0 @@
|
|||
# Braintrust Client Library
|
||||
|
||||
A Rust client library for interacting with the [Braintrust](https://braintrust.dev) API, allowing for logging of spans and traces to track AI application performance.
|
||||
|
||||
## Features
|
||||
|
||||
- Asynchronous logging of spans to Braintrust
|
||||
- Hierarchical tracing with parent-child relationships
|
||||
- Automatic timing and duration calculation
|
||||
- Token counting for LLM calls
|
||||
- Custom metadata and attributes
|
||||
- Background logging with non-blocking API
|
||||
|
||||
## Usage
|
||||
|
||||
Add the library to your project's dependencies:
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
braintrust = { path = "../libs/braintrust" }
|
||||
```
|
||||
|
||||
### Basic Example
|
||||
|
||||
```rust
|
||||
use braintrust::{BraintrustClient, TraceBuilder};
|
||||
use serde_json::json;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Initialize the client with your API key and project ID
|
||||
let client = BraintrustClient::new("YOUR_API_KEY", "YOUR_PROJECT_ID");
|
||||
|
||||
// Create a trace for the entire process
|
||||
let trace = TraceBuilder::new(client.clone(), "My AI Application");
|
||||
|
||||
// Create a span for an LLM call
|
||||
let mut span = trace.add_span("GPT-4 Call", "llm").await?;
|
||||
|
||||
// Update the span with input and output
|
||||
span = span
|
||||
.set_input(json!({
|
||||
"messages": [{"role": "user", "content": "Hello, world!"}]
|
||||
}))
|
||||
.set_output(json!({
|
||||
"choices": [{"message": {"role": "assistant", "content": "Hi there!"}}]
|
||||
}))
|
||||
.set_tokens(10, 5)
|
||||
.add_metadata("model", "gpt-4");
|
||||
|
||||
// Log the updated span
|
||||
client.log_span(span).await?;
|
||||
|
||||
// Finish the trace
|
||||
trace.finish().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Creating Spans
|
||||
|
||||
Spans represent a single operation or step in your application. You can create spans directly or through a `TraceBuilder`:
|
||||
|
||||
```rust
|
||||
// Create a span directly
|
||||
let span = Span::new("My Operation", "function", "root_span_id", Some("parent_span_id"));
|
||||
|
||||
// Or create a span through a trace builder
|
||||
let span = trace.add_span("My Operation", "function").await?;
|
||||
|
||||
// Create a child span with a parent reference
|
||||
let child_span = trace.add_child_span("Child Operation", "function", &parent_span).await?;
|
||||
```
|
||||
|
||||
### Updating Spans
|
||||
|
||||
Spans are immutable, but methods return a new span with updated fields:
|
||||
|
||||
```rust
|
||||
let updated_span = span
|
||||
.set_input(json!({ "key": "value" }))
|
||||
.set_output(json!({ "result": "success" }))
|
||||
.set_tokens(50, 30)
|
||||
.add_metadata("important_context", "some value");
|
||||
```
|
||||
|
||||
### Logging Spans
|
||||
|
||||
Spans can be logged asynchronously (non-blocking) or synchronously:
|
||||
|
||||
```rust
|
||||
// Asynchronous logging (non-blocking)
|
||||
client.log_span(span).await?;
|
||||
|
||||
// Synchronous logging (waits for API response)
|
||||
client.log_span_sync(span).await?;
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
See the examples directory for more advanced usage patterns:
|
||||
|
||||
- `basic_usage.rs`: Simple example of logging spans and traces
|
||||
- `conversation_tracking.rs`: Example of tracking a conversation with LLM calls
|
||||
|
||||
## Testing
|
||||
|
||||
Run the tests with:
|
||||
|
||||
```bash
|
||||
cargo test -p braintrust
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
|
@ -0,0 +1,113 @@
|
|||
# Database Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Database library serves as the core data layer for the Buster API, providing models, connection pooling, schema definitions, and database utilities. It acts as the foundation for data persistence and retrieval throughout the application.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Database connection pooling and management
|
||||
- Data models and schema definitions
|
||||
- Enum definitions for database-backed types
|
||||
- Helper functions for common database operations
|
||||
- Secure credential storage via the vault module
|
||||
- Structured types for complex data representations
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── enums.rs - Database-backed enum definitions
|
||||
├── helpers/
|
||||
│ ├── collections.rs - Collection-related helpers
|
||||
│ ├── dashboard_files.rs - Dashboard file operations
|
||||
│ ├── metric_files.rs - Metric file operations
|
||||
│ └── mod.rs
|
||||
├── lib.rs - Public exports
|
||||
├── models.rs - Database models and table definitions
|
||||
├── pool.rs - Database connection pooling
|
||||
├── schema.rs - Database schema definitions
|
||||
├── types/
|
||||
│ ├── dashboard_yml.rs - Dashboard YAML structure
|
||||
│ ├── metric_yml.rs - Metric YAML structure
|
||||
│ ├── mod.rs
|
||||
│ └── version_history.rs - Version tracking types
|
||||
└── vault.rs - Secure credential storage
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `models`: Contains the database model definitions that map to tables
|
||||
- `pool`: Manages database connections and connection pooling
|
||||
- `schema`: Defines the database schema structure
|
||||
- `enums`: Defines enum types used across the database
|
||||
- `vault`: Provides secure storage and retrieval of sensitive credentials
|
||||
- `helpers`: Contains utility functions for common database operations
|
||||
- `types`: Defines structured types for complex data
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use database::{pool::DbPool, models::User};
|
||||
use diesel::prelude::*;
|
||||
|
||||
async fn example_query(pool: &DbPool) -> Result<Vec<User>, anyhow::Error> {
|
||||
use database::schema::users::dsl::*;
|
||||
|
||||
// Get a connection from the pool
|
||||
let mut conn = pool.get().await?;
|
||||
|
||||
// Execute a query
|
||||
let results = users
|
||||
.filter(active.eq(true))
|
||||
.limit(10)
|
||||
.load::<User>(&mut conn)
|
||||
.await?;
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Always use the connection pool for database operations
|
||||
- Leverage Diesel's query builder for type-safe queries
|
||||
- Use the defined models and schema for table operations
|
||||
- Keep database operations isolated in repository-pattern functions
|
||||
- Use transactions for operations that modify multiple tables
|
||||
- Properly handle database errors and connection issues
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- None - this is a foundational library that other libraries depend on
|
||||
|
||||
- **External Dependencies**:
|
||||
- `diesel`: ORM for database operations
|
||||
- `diesel-async`: Async support for Diesel
|
||||
- `tokio-postgres`: PostgreSQL driver
|
||||
- `bb8-redis`: Redis connection pooling
|
||||
- `serde`: Serialization and deserialization
|
||||
- `chrono`: Date and time handling
|
||||
- `uuid`: UUID generation and handling
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `models.rs` to understand the database tables and their relationships
|
||||
- `schema.rs` defines the database schema and is generated from migrations
|
||||
- `pool.rs` contains the connection pool setup and management
|
||||
- Helper functions in the `helpers/` directory provide common database operations
|
||||
- When working with sensitive data, check the `vault.rs` module
|
||||
- Type definitions in the `types/` directory provide structure for complex data
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Use an in-memory or test database for unit tests
|
||||
- Create test fixtures for common database objects
|
||||
- Test database operations in isolation
|
||||
- Mock the database for non-database focused tests
|
||||
- Ensure proper cleanup of test data
|
||||
- Run tests with: `cargo test -p database`
|
||||
- Use transactions to roll back changes in tests
|
|
@ -0,0 +1,123 @@
|
|||
# Handlers Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Handlers library implements the request handlers for the Buster API, processing incoming HTTP requests, performing business logic, and returning appropriate responses. It acts as the controller layer in the application architecture, orchestrating interactions between other libraries.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- REST API endpoint handlers for all resources
|
||||
- Request validation and authentication
|
||||
- Response formatting and error handling
|
||||
- Business logic implementation
|
||||
- Resource management (chats, collections, dashboards, metrics, etc.)
|
||||
- Sharing and permission management
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── chats/ - Chat and thread handlers
|
||||
│ ├── context_loaders/ - Context loading for chats
|
||||
│ ├── helpers/ - Chat utility functions
|
||||
│ ├── sharing/ - Chat sharing functionality
|
||||
│ ├── streaming_parser.rs - Streaming response parsing
|
||||
│ ├── types.rs - Chat-specific types
|
||||
│ └── mod.rs
|
||||
├── collections/ - Collection management
|
||||
│ ├── sharing/ - Collection sharing
|
||||
│ ├── types.rs - Collection-specific types
|
||||
│ └── mod.rs
|
||||
├── dashboards/ - Dashboard management
|
||||
│ ├── sharing/ - Dashboard sharing
|
||||
│ ├── types.rs - Dashboard-specific types
|
||||
│ └── mod.rs
|
||||
├── data_sources/ - Data source handlers
|
||||
├── favorites/ - User favorites
|
||||
├── logs/ - Log handling
|
||||
├── messages/ - Message management
|
||||
│ ├── types/ - Message-related types
|
||||
│ └── helpers/ - Message utility functions
|
||||
├── metrics/ - Metrics management
|
||||
│ ├── sharing/ - Metrics sharing
|
||||
│ ├── types.rs - Metrics-specific types
|
||||
│ └── mod.rs
|
||||
├── utils/ - Shared utilities
|
||||
│ └── user/ - User-related utilities
|
||||
└── lib.rs - Public exports
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `chats`: Handles chat and thread operations (create, update, delete, list)
|
||||
- `collections`: Manages collections of assets (create, add/remove assets, share)
|
||||
- `dashboards`: Handles dashboard operations and associations
|
||||
- `data_sources`: Manages data source connections and configurations
|
||||
- `favorites`: Handles user favorite assets and resources
|
||||
- `messages`: Manages individual messages within threads
|
||||
- `metrics`: Handles metric operations and associations
|
||||
- `utils`: Shared utility functions used across handlers
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use handlers::chats::post_chat_handler;
|
||||
use handlers::thread_types::PostChatRequest;
|
||||
|
||||
async fn example_handler(req: PostChatRequest, user: AuthenticatedUser) -> Result<Json<Response>, Error> {
|
||||
// Call the appropriate handler
|
||||
let response = post_chat_handler(req, user, pool).await?;
|
||||
|
||||
// Return the response
|
||||
Ok(Json(response))
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Each handler follows a consistent pattern for request processing
|
||||
- Authentication and permission checks are performed at the beginning
|
||||
- Database operations are isolated and use the connection pool
|
||||
- Errors are mapped to appropriate HTTP responses
|
||||
- Responses are structured according to API specifications
|
||||
- Sharing handlers follow a consistent pattern across resource types
|
||||
- Context loaders are used to fetch related data for complex operations
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- `database`: For data persistence and retrieval
|
||||
- `agents`: For agent-based processing
|
||||
- `litellm`: For LLM interactions
|
||||
- `query_engine`: For executing queries against data sources
|
||||
- `middleware`: For authentication and request processing
|
||||
- `sharing`: For permission and sharing functionality
|
||||
|
||||
- **External Dependencies**:
|
||||
- `diesel` and `diesel-async`: For database operations
|
||||
- `serde` and `serde_json`: For serialization/deserialization
|
||||
- `uuid`: For unique identifier handling
|
||||
- `chrono`: For date/time operations
|
||||
- `anyhow`: For error handling
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see all exported modules
|
||||
- Each module in the top level represents a resource type (chats, collections, etc.)
|
||||
- Handler files are named after their operation (e.g., `create_dashboard_handler.rs`)
|
||||
- Type definitions for requests and responses are in `types.rs` files
|
||||
- Sharing functionality is consistently implemented in `sharing/` submodules
|
||||
- Helper functions are organized in `helpers/` submodules
|
||||
- Look for common patterns across similar handlers
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Mock dependencies (database, agents, etc.) for unit tests
|
||||
- Test error handling and edge cases
|
||||
- Verify permissions and access controls
|
||||
- Test the full request-response cycle for integration tests
|
||||
- Use test fixtures for consistent test data
|
||||
- Run tests with: `cargo test -p handlers`
|
||||
- Create tests for each handler in the corresponding `tests/` directory
|
|
@ -0,0 +1,96 @@
|
|||
# LiteLLM Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The LiteLLM library provides a unified interface for interacting with various Large Language Model (LLM) providers. It handles API communication, message formatting, response parsing, and error handling for multiple LLM services, simplifying the integration of LLMs throughout the application.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Unified client interface for multiple LLM providers
|
||||
- Message formatting and standardization
|
||||
- Streaming support for real-time responses
|
||||
- Error handling and retry logic
|
||||
- Structured types for LLM requests and responses
|
||||
- Agent message handling for agent-based workflows
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── client.rs - LLM client implementation
|
||||
├── types.rs - Types for LLM interaction
|
||||
└── lib.rs - Public exports
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `client`: Implements the LLM client that communicates with provider APIs
|
||||
- `types`: Defines structured types for requests, responses, and configurations
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use litellm::{LiteLLMClient, Message, Role};
|
||||
|
||||
async fn example_llm_call() -> Result<String, anyhow::Error> {
|
||||
// Create a client
|
||||
let client = LiteLLMClient::new("YOUR_API_KEY").await?;
|
||||
|
||||
// Prepare messages
|
||||
let messages = vec![
|
||||
Message {
|
||||
role: Role::System,
|
||||
content: "You are a helpful assistant.".to_string(),
|
||||
},
|
||||
Message {
|
||||
role: Role::User,
|
||||
content: "What is the capital of France?".to_string(),
|
||||
},
|
||||
];
|
||||
|
||||
// Send request to LLM
|
||||
let response = client.generate(messages, None).await?;
|
||||
|
||||
Ok(response.choices[0].message.content.clone())
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Initialize the client once and reuse it
|
||||
- Structure conversations as sequences of messages
|
||||
- Use system messages to set context and behavior
|
||||
- Handle streaming responses with callbacks
|
||||
- Include proper error handling for API failures
|
||||
- Set appropriate parameters for different use cases
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- None - this is a foundational library that other libraries depend on
|
||||
|
||||
- **External Dependencies**:
|
||||
- `reqwest`: For making HTTP requests to LLM providers
|
||||
- `serde_json`: For serializing and deserializing JSON
|
||||
- `tokio`: For async runtime support
|
||||
- `async-trait`: For async trait implementations
|
||||
- `futures` and `futures-util`: For async stream processing
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see what's exported
|
||||
- `client.rs` contains the main client implementation
|
||||
- `types.rs` defines the data structures for requests and responses
|
||||
- Look for provider-specific code in the client implementation
|
||||
- Understand the message structure and how it's processed
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Mock API responses using `mockito` for unit tests
|
||||
- Test different response formats and error conditions
|
||||
- Validate request formatting for each provider
|
||||
- Test streaming functionality with chunked responses
|
||||
- Run tests with: `cargo test -p litellm`
|
||||
- Use environment variables for API keys in integration tests
|
|
@ -0,0 +1,90 @@
|
|||
# Middleware Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Middleware library provides HTTP middleware components for the Buster API, handling cross-cutting concerns like authentication, authorization, and CORS. It serves as the security and request processing layer that runs before request handlers.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Authentication middleware for validating user credentials
|
||||
- User identity extraction and verification
|
||||
- CORS (Cross-Origin Resource Sharing) configuration
|
||||
- Request validation and preprocessing
|
||||
- Structured types for authentication data
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── auth.rs - Authentication middleware
|
||||
├── cors.rs - CORS configuration
|
||||
├── types.rs - Shared type definitions
|
||||
└── lib.rs - Public exports and documentation
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `auth`: Implements authentication middleware for validating tokens and extracting user information
|
||||
- `cors`: Configures CORS headers and policies for cross-origin requests
|
||||
- `types`: Defines shared types used across middleware components
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use axum::{Router, routing::get};
|
||||
use middleware::{auth, cors, types::AuthenticatedUser};
|
||||
|
||||
async fn create_router() -> Router {
|
||||
Router::new()
|
||||
.route("/protected", get(protected_handler))
|
||||
.layer(auth()) // Add authentication middleware
|
||||
.layer(cors()) // Add CORS middleware
|
||||
}
|
||||
|
||||
async fn protected_handler(user: AuthenticatedUser) -> impl IntoResponse {
|
||||
// This handler only runs if authentication succeeds
|
||||
// The user parameter contains the authenticated user information
|
||||
format!("Hello, {}!", user.email)
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Apply middleware at the router level for global concerns
|
||||
- Extract authenticated user information in handlers that require it
|
||||
- Use organization and team membership for authorization checks
|
||||
- Handle authentication errors with appropriate status codes
|
||||
- Configure CORS for allowed origins and methods
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- `database`: For user lookup and validation
|
||||
|
||||
- **External Dependencies**:
|
||||
- `axum`: Web framework integration
|
||||
- `tower-http`: HTTP middleware components
|
||||
- `jsonwebtoken`: JWT validation and parsing
|
||||
- `serde`: Serialization and deserialization
|
||||
- `uuid`: For user and organization IDs
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see the exported functions and types
|
||||
- `auth.rs` contains the authentication middleware implementation
|
||||
- `cors.rs` contains the CORS configuration
|
||||
- `types.rs` defines important types like `AuthenticatedUser`
|
||||
- Look for the `auth()` function as the main entry point for authentication
|
||||
- Understand the structure of `AuthenticatedUser` to see what information is available
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Test authentication with valid and invalid tokens
|
||||
- Verify that protected routes reject unauthenticated requests
|
||||
- Test CORS with different origin configurations
|
||||
- Mock database responses for authentication tests
|
||||
- Use test JWT tokens for authentication testing
|
||||
- Run tests with: `cargo test -p middleware`
|
||||
- Test integration with axum route handlers
|
|
@ -0,0 +1,113 @@
|
|||
# Query Engine Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Query Engine library provides connectivity and query execution functionality for various data sources in Buster. It abstracts away the details of connecting to different database systems, allows secure credential management, and provides a unified interface for executing queries across multiple database technologies.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Data source connection management for multiple database types
|
||||
- Secure credential handling and storage
|
||||
- Query execution across various database systems
|
||||
- Data type conversions and mapping
|
||||
- Connection testing and validation
|
||||
- SSH tunneling for secure remote connections
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── credentials.rs - Credential management
|
||||
├── data_source_connections/ - Database connectors
|
||||
│ ├── get_bigquery_client.rs - BigQuery connection
|
||||
│ ├── get_databricks_client.rs - Databricks connection
|
||||
│ ├── get_mysql_connection.rs - MySQL connection
|
||||
│ ├── get_postgres_connection.rs - PostgreSQL connection
|
||||
│ ├── get_redshift_connection.rs - Redshift connection
|
||||
│ ├── get_snowflake_client.rs - Snowflake connection
|
||||
│ ├── get_sql_server_connection.rs - SQL Server connection
|
||||
│ ├── ssh_tunneling.rs - SSH tunnel functionality
|
||||
│ ├── test_data_source_connections.rs - Connection testing
|
||||
│ └── mod.rs
|
||||
├── data_source_helpers.rs - Helper functions
|
||||
├── data_source_query_routes/ - Query executors
|
||||
│ ├── bigquery_query.rs - BigQuery execution
|
||||
│ ├── databricks_query.rs - Databricks execution
|
||||
│ ├── mysql_query.rs - MySQL execution
|
||||
│ ├── postgres_query.rs - PostgreSQL execution
|
||||
│ ├── query_engine.rs - Core query engine
|
||||
│ ├── redshift_query.rs - Redshift execution
|
||||
│ ├── security_utils.rs - Security utilities
|
||||
│ ├── snowflake_query.rs - Snowflake execution
|
||||
│ ├── sql_server_query.rs - SQL Server execution
|
||||
│ └── mod.rs
|
||||
├── data_types.rs - Common data types
|
||||
└── lib.rs - Public exports
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `data_source_connections`: Database connection handlers for different database systems
|
||||
- `data_source_query_routes`: Query execution implementations for each database system
|
||||
- `credentials`: Secure credential management for database connections
|
||||
- `data_types`: Common data type definitions and conversions
|
||||
- `data_source_helpers`: Utility functions for data sources
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use query_engine::data_source_query_routes::query_engine::{execute_query, QueryResult};
|
||||
use query_engine::data_types::{DataSource, DataSourceType};
|
||||
|
||||
async fn example_query(data_source: DataSource) -> Result<QueryResult, anyhow::Error> {
|
||||
// Execute a query against a data source
|
||||
let query = "SELECT * FROM users LIMIT 10";
|
||||
let result = execute_query(&data_source, query).await?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Use the appropriate connection function for each database type
|
||||
- Handle connection pooling and reuse when possible
|
||||
- Secure credentials through the credentials module
|
||||
- Use SSH tunneling for remote databases that require it
|
||||
- Handle database-specific query syntax differences
|
||||
- Convert results to a common format for the API
|
||||
- Implement proper error handling for connection and query failures
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- `database`: For storing and retrieving data source configurations
|
||||
|
||||
- **External Dependencies**:
|
||||
- Database-specific clients: `sqlx`, `gcp-bigquery-client`, `snowflake-api`, `tiberius`
|
||||
- `tokio`: For async runtime
|
||||
- `anyhow`: For error handling
|
||||
- `serde` and `serde_json`: For serialization
|
||||
- `uuid`: For unique identifiers
|
||||
- `arrow`: For columnar data processing
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see what's exported
|
||||
- The central component is the `execute_query` function in `data_source_query_routes/query_engine.rs`
|
||||
- Connection functions in `data_source_connections/` follow a naming pattern of `get_<database>_connection`
|
||||
- Query execution in `data_source_query_routes/` follows a similar pattern
|
||||
- Database-specific code is isolated in dedicated modules
|
||||
- Common data types are defined in `data_types.rs`
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Use mock connections for unit testing
|
||||
- Test each database connector separately
|
||||
- Verify proper error handling for connection failures
|
||||
- Test query execution with various input types
|
||||
- Test data type conversions for each database system
|
||||
- Run tests with: `cargo test -p query_engine`
|
||||
- Use test fixtures for database configurations
|
|
@ -0,0 +1,132 @@
|
|||
# Sharing Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Sharing library provides functionality for managing asset permissions and sharing resources with users and teams in Buster. It implements the core sharing model, handling permission checks, creation, listing, and removal of access rights across all asset types.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Permission checking for asset access
|
||||
- Creation of sharing permissions
|
||||
- Listing existing permissions on assets
|
||||
- Removal of sharing permissions
|
||||
- User lookup by email for sharing
|
||||
- Error handling specific to sharing operations
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── check_asset_permission.rs - Permission verification
|
||||
├── create_asset_permission.rs - Creating new permissions
|
||||
├── errors.rs - Sharing-specific error types
|
||||
├── list_asset_permissions.rs - Listing existing permissions
|
||||
├── remove_asset_permissions.rs - Removing permissions
|
||||
├── types.rs - Sharing-related type definitions
|
||||
├── user_lookup.rs - Finding users for sharing
|
||||
├── tests/ - Test modules
|
||||
│ ├── check_asset_permission_test.rs
|
||||
│ └── mod.rs
|
||||
└── lib.rs - Public exports and documentation
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `check_asset_permission`: Functions for verifying if a user has access to assets
|
||||
- `create_asset_permission`: Functions for granting access to assets
|
||||
- `list_asset_permissions`: Functions for listing existing permissions
|
||||
- `remove_asset_permissions`: Functions for revoking access to assets
|
||||
- `user_lookup`: Functions for finding users by email for sharing
|
||||
- `types`: Data structures for sharing operations
|
||||
- `errors`: Error types specific to sharing operations
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use sharing::{check_access, create_share_by_email, list_shares, remove_share_by_email};
|
||||
use sharing::types::{AssetType, Permission};
|
||||
use uuid::Uuid;
|
||||
|
||||
async fn example_sharing_operations(pool: &DbPool) -> Result<(), anyhow::Error> {
|
||||
let asset_id = Uuid::new_v4();
|
||||
let asset_type = AssetType::Dashboard;
|
||||
let user_email = "user@example.com";
|
||||
|
||||
// Create a new share
|
||||
create_share_by_email(
|
||||
pool,
|
||||
asset_id,
|
||||
asset_type,
|
||||
user_email,
|
||||
Permission::View,
|
||||
None, // organization_id (if applicable)
|
||||
).await?;
|
||||
|
||||
// Check if user has access
|
||||
let has_access = check_access(
|
||||
pool,
|
||||
asset_id,
|
||||
asset_type,
|
||||
Some(user_id), // User ID to check
|
||||
Permission::View,
|
||||
).await?;
|
||||
|
||||
// List all shares for an asset
|
||||
let shares = list_shares(
|
||||
pool,
|
||||
asset_id,
|
||||
asset_type,
|
||||
None, // organization_id (if applicable)
|
||||
).await?;
|
||||
|
||||
// Remove a share
|
||||
remove_share_by_email(
|
||||
pool,
|
||||
asset_id,
|
||||
asset_type,
|
||||
user_email,
|
||||
).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Always check permissions before allowing access to restricted operations
|
||||
- Handle not-found errors appropriately when users don't exist
|
||||
- Use bulk operations when dealing with multiple assets or users
|
||||
- Verify organization context for multi-tenant operations
|
||||
- Use appropriate permission levels (View, Edit, Admin) based on operation
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- `database`: For storing and retrieving permission records
|
||||
|
||||
- **External Dependencies**:
|
||||
- `diesel` and `diesel-async`: For database operations
|
||||
- `uuid`: For working with unique identifiers
|
||||
- `thiserror`: For error handling
|
||||
- `chrono`: For timestamp handling
|
||||
- `serde`: For serialization of sharing data
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see the exported functions
|
||||
- Key operations are in separate modules with descriptive names
|
||||
- `check_asset_permission.rs` contains the core permission checking logic
|
||||
- `types.rs` defines the sharing model data structures
|
||||
- Error handling is centralized in `errors.rs`
|
||||
- Database queries are typically implemented in the module corresponding to their operation
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Test permission checks with various access levels
|
||||
- Verify that permissions are correctly created, listed, and removed
|
||||
- Test error cases like sharing with non-existent users
|
||||
- Use test fixtures for consistent test data
|
||||
- Run tests with: `cargo test -p sharing`
|
||||
- Check test coverage for critical permission checks
|
|
@ -0,0 +1,113 @@
|
|||
# Streaming Library - Agent Guidance
|
||||
|
||||
## Purpose & Role
|
||||
|
||||
The Streaming library provides functionality for parsing and processing incomplete JSON streams from LLM responses. It allows for real-time processing of partial JSON data, handling tool calls and message content as they arrive, and supports concurrent processing of multiple tool calls of the same type.
|
||||
|
||||
## Key Functionality
|
||||
|
||||
- Streaming JSON parser for incomplete LLM responses
|
||||
- Tool call detection and processing
|
||||
- Processor registry for handling different types of tool calls
|
||||
- State tracking for partial message processing
|
||||
- Caching mechanism for tool call chunks
|
||||
- Specialized processors for common tools
|
||||
|
||||
## Internal Organization
|
||||
|
||||
### Directory Structure
|
||||
|
||||
```
|
||||
src/
|
||||
├── parser.rs - StreamingParser implementation
|
||||
├── processor.rs - Processor trait and registry
|
||||
├── types.rs - Core type definitions
|
||||
├── processors/ - Specialized processors
|
||||
│ ├── create_dashboards_processor.rs - Dashboard creation
|
||||
│ ├── create_metrics_processor.rs - Metric creation
|
||||
│ ├── create_plan_processor.rs - Plan creation
|
||||
│ ├── search_data_catalog_processor.rs - Data catalog search
|
||||
│ └── mod.rs
|
||||
└── lib.rs - Public exports and documentation
|
||||
```
|
||||
|
||||
### Key Modules
|
||||
|
||||
- `parser`: Implements the `StreamingParser` that processes JSON chunks
|
||||
- `processor`: Defines the `Processor` trait and the `ProcessorRegistry`
|
||||
- `types`: Contains core types like `ProcessedOutput` and `ToolCallInfo`
|
||||
- `processors`: Specialized implementations for specific tool calls
|
||||
|
||||
## Usage Patterns
|
||||
|
||||
```rust
|
||||
use streaming::{StreamingParser, Processor, ProcessorRegistry, ToolCallInfo, ProcessedOutput};
|
||||
|
||||
async fn example_streaming(chunk: String, registry: &ProcessorRegistry) -> Result<Vec<ProcessedOutput>, anyhow::Error> {
|
||||
// Create a parser
|
||||
let mut parser = StreamingParser::new();
|
||||
|
||||
// Process a chunk
|
||||
let outputs = parser.process_chunk(&chunk, registry).await?;
|
||||
|
||||
// Handle the processed outputs
|
||||
for output in &outputs {
|
||||
match output {
|
||||
ProcessedOutput::ToolStart(tool_info) => {
|
||||
println!("Tool started: {}", tool_info.name);
|
||||
},
|
||||
ProcessedOutput::ToolOutput(tool_id, output) => {
|
||||
println!("Tool output for {}: {}", tool_id, output);
|
||||
},
|
||||
ProcessedOutput::ToolEnd(tool_id) => {
|
||||
println!("Tool completed: {}", tool_id);
|
||||
},
|
||||
ProcessedOutput::Message(content) => {
|
||||
println!("Message content: {}", content);
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
Ok(outputs)
|
||||
}
|
||||
```
|
||||
|
||||
### Common Implementation Patterns
|
||||
|
||||
- Register processors before parsing begins
|
||||
- Process chunks sequentially as they arrive
|
||||
- Handle both message content and tool calls
|
||||
- Use the processor registry to delegate tool call processing
|
||||
- Track the state of partial tool calls across chunks
|
||||
- Cache partial chunks until complete JSON objects are received
|
||||
|
||||
## Dependencies
|
||||
|
||||
- **Internal Dependencies**:
|
||||
- `litellm`: For LLM message types and structures
|
||||
|
||||
- **External Dependencies**:
|
||||
- `serde_json`: For JSON parsing
|
||||
- `regex`: For pattern matching in JSON chunks
|
||||
- `uuid`: For unique identifiers
|
||||
- `sha2`: For hash generation
|
||||
- `chrono`: For timestamps
|
||||
|
||||
## Code Navigation Tips
|
||||
|
||||
- Start with `lib.rs` to see the exported types and functions
|
||||
- `StreamingParser` in `parser.rs` is the main entry point
|
||||
- `Processor` trait in `processor.rs` defines the interface for tool processors
|
||||
- `ProcessorRegistry` manages the registered processors
|
||||
- Each processor in the `processors/` directory implements the `Processor` trait
|
||||
- `types.rs` defines the core data structures used throughout the library
|
||||
|
||||
## Testing Guidelines
|
||||
|
||||
- Test with partial and complete JSON chunks
|
||||
- Verify correct handling of malformed JSON
|
||||
- Test concurrent processing of multiple tool calls
|
||||
- Test each processor individually
|
||||
- Test the full streaming pipeline with realistic LLM responses
|
||||
- Run tests with: `cargo test -p streaming`
|
||||
- Create test fixtures for common LLM response patterns
|
|
@ -1,118 +0,0 @@
|
|||
# Update Message REST Endpoint PRD
|
||||
|
||||
## Overview
|
||||
This PRD describes the implementation of a new REST endpoint for updating a message. This endpoint will allow users to update specific properties of a message, such as setting feedback on a message.
|
||||
|
||||
## Endpoint Details
|
||||
|
||||
### Basic Information
|
||||
- **HTTP Method**: PUT
|
||||
- **URL Path**: `/api/v1/messages/:id`
|
||||
- **Purpose**: Update properties of a specific message
|
||||
|
||||
### Request Format
|
||||
- **Headers**:
|
||||
- `Content-Type: application/json`
|
||||
- `Authorization: Bearer {token}`
|
||||
- **URL Parameters**:
|
||||
- `id`: UUID of the message to update
|
||||
- **Request Body (JSON)**:
|
||||
- `feedback`: String (optional) - "positive" or "negative"
|
||||
|
||||
### Response Format
|
||||
- **Success**:
|
||||
- Status: 200 OK
|
||||
- Body: Updated message object
|
||||
- **Error Responses**:
|
||||
- 401 Unauthorized: If the user is not authenticated
|
||||
- 403 Forbidden: If the user does not have permission to update the message
|
||||
- 404 Not Found: If the message with the specified ID does not exist
|
||||
- 500 Internal Server Error: If there's a server error
|
||||
|
||||
### Example Request
|
||||
```bash
|
||||
curl --location --request PUT 'http://localhost:3001/api/v1/messages/cee3065d-c165-4e04-be44-be05d1e6d902' \
|
||||
--header 'Content-Type: application/json' \
|
||||
--header 'Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJjMmRkNjRjZC1mN2YzLTQ4ODQtYmM5MS1kNDZhZTQzMTkwMWUiLCJuYW1lIjoiSm9obiBEb2UiLCJpYXQiOjE1MTYyMzkwMjIsImV4cCI6MjUxODU1NTMyNiwiYXVkIjoiYXV0aGVudGljYXRlZCJ9.uRs5OVyYErQ1iSQwAXVUD6TFolOu31ejPhcBS41ResA' \
|
||||
--data '{
|
||||
"feedback": "negative"
|
||||
}'
|
||||
```
|
||||
|
||||
### Example Response
|
||||
```json
|
||||
{
|
||||
"id": "cee3065d-c165-4e04-be44-be05d1e6d902",
|
||||
"request_message": {
|
||||
"request": "Example message text",
|
||||
"sender_id": "c2dd64cd-f7f3-4884-bc91-d46ae431901e",
|
||||
"sender_name": "John Doe",
|
||||
"sender_avatar": null
|
||||
},
|
||||
"response_message_ids": ["..."],
|
||||
"response_messages": {},
|
||||
"reasoning_message_ids": ["..."],
|
||||
"reasoning_messages": {},
|
||||
"created_at": "2025-03-20T15:30:00Z",
|
||||
"final_reasoning_message": null,
|
||||
"feedback": "negative"
|
||||
}
|
||||
```
|
||||
|
||||
## Implementation Details
|
||||
|
||||
### New Files to Create
|
||||
|
||||
#### 1. Handler Implementation
|
||||
- **File**: `/libs/handlers/src/messages/helpers/update_message_handler.rs`
|
||||
- **Purpose**: Contains the business logic for updating a message
|
||||
|
||||
#### 2. REST Route Implementation
|
||||
- **File**: `/src/routes/rest/routes/messages/update_message.rs`
|
||||
- **Purpose**: Contains the HTTP handler for the REST endpoint
|
||||
|
||||
#### 3. Update Module Exports
|
||||
- **File**: `/libs/handlers/src/messages/helpers/mod.rs` (update)
|
||||
- **File**: `/src/routes/rest/routes/messages/mod.rs` (update)
|
||||
|
||||
### Implementation Steps
|
||||
|
||||
#### 1. Create the Handler Function
|
||||
Create a new handler in `/libs/handlers/src/messages/helpers/update_message_handler.rs` that:
|
||||
- Takes a user, message ID, and update parameters
|
||||
- Verifies the user has permission to update the message
|
||||
- Updates the message in the database
|
||||
- Returns the updated message
|
||||
|
||||
#### 2. Create the REST Endpoint
|
||||
Create a new file at `/src/routes/rest/routes/messages/update_message.rs` that:
|
||||
- Creates a request struct to deserialize the request body
|
||||
- Maps the HTTP request to the handler function
|
||||
- Handles errors appropriately
|
||||
- Returns the appropriate HTTP response
|
||||
|
||||
#### 3. Update Module Exports
|
||||
Update the module exports to include the new handler and endpoint.
|
||||
|
||||
## Technical Details
|
||||
|
||||
### Database Updates
|
||||
The update operation will modify records in the `messages` table.
|
||||
|
||||
### Expected Behavior
|
||||
- Only the user who created the message or an admin should be able to update the message
|
||||
- Updates should be atomic and consistent
|
||||
- The endpoint should validate input data before updating the database
|
||||
|
||||
### Testing Strategy
|
||||
1. Test updating a message with valid data as the message owner
|
||||
2. Test updating a message with valid data as an admin
|
||||
3. Test updating a message with invalid data
|
||||
4. Test updating a non-existent message
|
||||
5. Test updating a message without permission
|
||||
6. Test concurrent updates to the same message
|
||||
|
||||
## Security Considerations
|
||||
- Ensure proper authentication and authorization
|
||||
- Validate input data to prevent injection attacks
|
||||
- Ensure database queries are properly parameterized
|
|
@ -12,7 +12,7 @@ pub async fn update_data_source(
|
|||
Json(payload): Json<UpdateDataSourceRequest>,
|
||||
) -> Result<ApiResponse<()>, (StatusCode, &'static str)> {
|
||||
match update_data_source_handler(&user.id, &id, payload).await {
|
||||
Ok(data_source) => Ok(ApiResponse::NoContent),
|
||||
Ok(_) => Ok(ApiResponse::NoContent),
|
||||
Err(e) => {
|
||||
tracing::error!("Error updating data source: {:?}", e);
|
||||
Err((
|
||||
|
|
Loading…
Reference in New Issue