Compare commits
83 Commits
v1.0.14-al
...
v2.3.1-alf
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
67078ce925 | ||
|
|
ebdb836448 | ||
|
|
81e754317a | ||
|
|
578981c745 | ||
|
|
8fb2ad43c5 | ||
|
|
49f9077a7b | ||
|
|
d290b46a0c | ||
|
|
73647e4795 | ||
|
|
25e169dbea | ||
|
|
8a29eb0d8f | ||
|
|
0a5f0986e6 | ||
|
|
4d79c4fd5a | ||
|
|
5123de55cc | ||
|
|
1fdbd2ff45 | ||
|
|
d789e431ca | ||
|
|
70de4c0328 | ||
|
|
d2bb51a4a8 | ||
|
|
28aea85b10 | ||
|
|
d2a9092f46 | ||
|
|
5c982fcc2c | ||
|
|
b4f7b210e0 | ||
|
|
1b1eef0d2e | ||
|
|
17d32cd039 | ||
|
|
12a53ebc1c | ||
|
|
a421977918 | ||
|
|
4c480c9baa | ||
|
|
9ea04572c8 | ||
|
|
6ef025363d | ||
|
|
9652d0bff9 | ||
|
|
4bf12db142 | ||
|
|
5f58417d24 | ||
|
|
3eed546879 | ||
|
|
35f0adef1b | ||
|
|
f43e79376c | ||
|
|
be76dd5240 | ||
|
|
c2c3b01b28 | ||
|
|
8daa52d1e9 | ||
|
|
9ad7c1aee9 | ||
|
|
1762b930bc | ||
|
|
d57bc5cf03 | ||
|
|
6c8c33d296 | ||
|
|
4ea16521e2 | ||
|
|
b6ee7182de | ||
|
|
238bdb58f4 | ||
|
|
a35486b573 | ||
|
|
dc64bbc257 | ||
|
|
09555ae8b0 | ||
|
|
cf2201a1f7 | ||
|
|
a6402524ce | ||
|
|
56a00c2894 | ||
|
|
6465e4f358 | ||
|
|
4b43f96afe | ||
|
|
e088ef7e4e | ||
|
|
9e03af45e1 | ||
|
|
5bfd3445bb | ||
|
|
efff63043a | ||
|
|
c15cabc289 | ||
|
|
55a89c11bb | ||
|
|
c037d4135e | ||
|
|
25213f2004 | ||
|
|
d106520d22 | ||
|
|
7bddeb0ebd | ||
|
|
f7cd58ed2a | ||
|
|
53c625599a | ||
|
|
88ee4f482b | ||
|
|
3176b95323 | ||
|
|
46c60b36a0 | ||
|
|
d35ec9f5ae | ||
|
|
311927d5ea | ||
|
|
fb798501b9 | ||
|
|
99135c9b02 | ||
|
|
425b580f15 | ||
|
|
b658e68e65 | ||
|
|
b8e07bec77 | ||
|
|
344ea26ecc | ||
|
|
98cb4e4f2f | ||
|
|
07d89d204f | ||
|
|
7702a6dfcc | ||
|
|
4c009949b3 | ||
|
|
aa4ac3ec7c | ||
|
|
1807435339 | ||
|
|
55a8a95f79 | ||
|
|
503ea7965d |
12
.gitignore
vendored
12
.gitignore
vendored
@@ -14,7 +14,6 @@ __pycache__
|
||||
**/__pycache__
|
||||
/.idea
|
||||
*.pyc
|
||||
*.pyc
|
||||
common/.DS_Store
|
||||
common/__pycache__/__init__.cpython-312.pyc
|
||||
common/__pycache__/extensions.cpython-312.pyc
|
||||
@@ -43,3 +42,14 @@ scripts/.DS_Store
|
||||
scripts/__pycache__/run_eveai_app.cpython-312.pyc
|
||||
/eveai_repo.txt
|
||||
*repo.txt
|
||||
/docker/eveai_logs/
|
||||
/integrations/Wordpress/eveai_sync.zip
|
||||
/integrations/Wordpress/eveai-chat.zip
|
||||
/db_backups/
|
||||
/tests/interactive_client/specialist_client.log
|
||||
/.repopackignore
|
||||
/patched_packages/crewai/
|
||||
/docker/prometheus/data/
|
||||
/docker/grafana/data/
|
||||
/temp_requirements/
|
||||
/nginx/node_modules/
|
||||
|
||||
8
.idea/.gitignore
generated
vendored
8
.idea/.gitignore
generated
vendored
@@ -1,8 +0,0 @@
|
||||
# Default ignored files
|
||||
/shelf/
|
||||
/workspace.xml
|
||||
# Editor-based HTTP Client requests
|
||||
/httpRequests/
|
||||
# Datasource local storage ignored files
|
||||
/dataSources/
|
||||
/dataSources.local.xml
|
||||
22
.idea/eveAI.iml
generated
22
.idea/eveAI.iml
generated
@@ -1,22 +0,0 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<module type="PYTHON_MODULE" version="4">
|
||||
<component name="Flask">
|
||||
<option name="enabled" value="true" />
|
||||
</component>
|
||||
<component name="NewModuleRootManager">
|
||||
<content url="file://$MODULE_DIR$">
|
||||
<excludeFolder url="file://$MODULE_DIR$/.venv" />
|
||||
<excludeFolder url="file://$MODULE_DIR$/.venv2" />
|
||||
</content>
|
||||
<orderEntry type="jdk" jdkName="Python 3.12 (eveai_dev)" jdkType="Python SDK" />
|
||||
<orderEntry type="sourceFolder" forTests="false" />
|
||||
</component>
|
||||
<component name="TemplatesService">
|
||||
<option name="TEMPLATE_CONFIGURATION" value="Jinja2" />
|
||||
<option name="TEMPLATE_FOLDERS">
|
||||
<list>
|
||||
<option value="$MODULE_DIR$/templates" />
|
||||
</list>
|
||||
</option>
|
||||
</component>
|
||||
</module>
|
||||
6
.idea/inspectionProfiles/profiles_settings.xml
generated
6
.idea/inspectionProfiles/profiles_settings.xml
generated
@@ -1,6 +0,0 @@
|
||||
<component name="InspectionProjectProfileManager">
|
||||
<settings>
|
||||
<option name="USE_PROJECT_PROFILE" value="false" />
|
||||
<version value="1.0" />
|
||||
</settings>
|
||||
</component>
|
||||
7
.idea/misc.xml
generated
7
.idea/misc.xml
generated
@@ -1,7 +0,0 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="Black">
|
||||
<option name="sdkName" value="Python 3.12 (eveai_tbd)" />
|
||||
</component>
|
||||
<component name="ProjectRootManager" version="2" project-jdk-name="Python 3.12 (eveai_tbd)" project-jdk-type="Python SDK" />
|
||||
</project>
|
||||
8
.idea/modules.xml
generated
8
.idea/modules.xml
generated
@@ -1,8 +0,0 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="ProjectModuleManager">
|
||||
<modules>
|
||||
<module fileurl="file://$PROJECT_DIR$/.idea/TBD.iml" filepath="$PROJECT_DIR$/.idea/TBD.iml" />
|
||||
</modules>
|
||||
</component>
|
||||
</project>
|
||||
6
.idea/vcs.xml
generated
6
.idea/vcs.xml
generated
@@ -1,6 +0,0 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project version="4">
|
||||
<component name="VcsDirectoryMappings">
|
||||
<mapping directory="$PROJECT_DIR$" vcs="Git" />
|
||||
</component>
|
||||
</project>
|
||||
@@ -1 +0,0 @@
|
||||
eveai_tbd
|
||||
156
CHANGELOG.md
156
CHANGELOG.md
@@ -5,10 +5,115 @@ All notable changes to EveAI will be documented in this file.
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
|
||||
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
|
||||
|
||||
## [Unreleased]
|
||||
## [2.3.0-alfa]
|
||||
|
||||
### Added
|
||||
- For new features.
|
||||
- Introduction of Push Gateway for Prometheus
|
||||
- Introduction of Partner Models
|
||||
- Introduction of Tenant and Partner codes for more security
|
||||
- Introduction of 'Management Partner' type and additional 'Partner Admin'-role
|
||||
- Introduction of a technical services layer
|
||||
- Introduction of partner-specific configurations
|
||||
- Introduction of additional test environment
|
||||
- Introduction of strict no-overage usage
|
||||
- Introduction of LicensePeriod, Payments & Invoices
|
||||
- Introduction of Processed File Viewer
|
||||
- Introduction of Traicie Role Definition Specialist
|
||||
- Allow invocation of non-interactive specialists in administrative interface (eveai_app)
|
||||
- Introduction of advanced JSON editor
|
||||
- Introduction of ChatSession (Specialist Execution) follow-up in administrative interface
|
||||
- Introduce npm for javascript libraries usage and optimisations
|
||||
- Introduction of new top bar in administrative interface to show session defaults (removing old navbar buttons)
|
||||
-
|
||||
|
||||
### Changed
|
||||
- Add 'Register'-button to list views, replacing register menu-items
|
||||
- Add additional environment capabilities in docker
|
||||
- PDF Processor now uses Mistral OCR
|
||||
- Allow additional chunking mechanisms for very long chunks (in case of very large documents)
|
||||
- Allow for TrackedMistralAIEmbedding batching to allow for processing long documents
|
||||
- RAG & SPIN Specialist improvements
|
||||
- Move mail messaging from standard SMTP to Scaleway TEM mails
|
||||
- Improve mail layouts
|
||||
- Add functionality to add a default dictionary for dynamic forms
|
||||
- AI model choices defined by Ask Eve AI iso Tenant (replaces ModelVariables completely)
|
||||
- Improve HTML Processing
|
||||
- Pagination improvements
|
||||
- Update Material Kit Pro to latest version
|
||||
|
||||
### Removed
|
||||
- Repopack implementation ==> Using PyCharm's new AI capabilities instead
|
||||
|
||||
### Fixed
|
||||
- Synchronous vs Asynchronous behaviour in crewAI type specialists
|
||||
- Nasty dynamic boolean fields bug corrected
|
||||
- Several smaller bugfixes
|
||||
- Tasks & Tools editors finished
|
||||
|
||||
### Security
|
||||
- In case of vulnerabilities.
|
||||
|
||||
## [2.2.0-alfa]
|
||||
|
||||
### Added
|
||||
- Mistral AI as main provider for embeddings, chains and specialists
|
||||
- Usage measuring for specialists
|
||||
- RAG from chain to specialist technology
|
||||
- Dossier catalog management possibilities added to eveai_app
|
||||
- Asset definition (Paused - other priorities)
|
||||
- Prometheus and Grafana
|
||||
- Add prometheus monitoring to business events
|
||||
- Asynchronous execution of specialists
|
||||
|
||||
### Changed
|
||||
- Moved choice for AI providers / models to specialists and prompts
|
||||
- Improve RAG to not repeat historic answers
|
||||
- Fixed embedding model, no more choices allowed
|
||||
- clean url (of tracking parameters) before adding it to a catalog
|
||||
|
||||
### Deprecated
|
||||
- For soon-to-be removed features.
|
||||
|
||||
### Removed
|
||||
- Add Multiple URLs removed from menu
|
||||
- Old Specialist items removed from interaction menu
|
||||
-
|
||||
|
||||
### Fixed
|
||||
- Set default language when registering Documents or URLs.
|
||||
|
||||
### Security
|
||||
- In case of vulnerabilities.
|
||||
|
||||
## [2.1.0-alfa]
|
||||
|
||||
### Added
|
||||
- Zapier Refresh Document
|
||||
- SPIN Specialist definition - from start to finish
|
||||
- Introduction of startup scripts in eveai_app
|
||||
- Caching for all configurations added
|
||||
- Caching for processed specialist configurations
|
||||
- Caching for specialist history
|
||||
- Augmented Specialist Editor, including Specialist graphic presentation
|
||||
- Introduction of specialist_execution_api, introducting SSE
|
||||
- Introduction of crewai framework for specialist implementation
|
||||
- Test app for testing specialists - also serves as a sample client application for SSE
|
||||
-
|
||||
|
||||
### Changed
|
||||
- Improvement of startup of applications using gevent, and better handling and scaling of multiple connections
|
||||
- STANDARD_RAG Specialist improvement
|
||||
-
|
||||
|
||||
### Deprecated
|
||||
- eveai_chat - using sockets - will be replaced with new specialist_execution_api and SSE
|
||||
|
||||
## [2.0.1-alfa]
|
||||
|
||||
### Added
|
||||
- Zapîer Integration (partial - only adding files).
|
||||
- Addition of general chunking parameters (chunking_heading_level and chunking_patterns)
|
||||
- Addition of DocX and markdown Processor Types
|
||||
|
||||
### Changed
|
||||
- For changes in existing functionality.
|
||||
@@ -19,6 +124,53 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
### Removed
|
||||
- For now removed features.
|
||||
|
||||
### Fixed
|
||||
- Ensure the RAG Specialist is using the detailed_question
|
||||
- Wordpress Chat Plugin: languages dropdown filled again
|
||||
- OpenAI update - proxies no longer supported
|
||||
- Build & Release script for Wordpress Plugins (including end user download folder)
|
||||
|
||||
### Security
|
||||
- In case of vulnerabilities.
|
||||
|
||||
## [2.0.0-alfa]
|
||||
|
||||
### Added
|
||||
- Introduction of dynamic Retrievers & Specialists
|
||||
- Introduction of dynamic Processors
|
||||
- Introduction of caching system
|
||||
- Introduction of a better template manager
|
||||
- Modernisation of external API/Socket authentication using projects
|
||||
- Creation of new eveai_chat WordPress plugin to support specialists
|
||||
|
||||
### Changed
|
||||
- Update of eveai_sync WordPress plugin
|
||||
|
||||
### Fixed
|
||||
- Set default language when registering Documents or URLs.
|
||||
|
||||
### Security
|
||||
- Security improvements to Docker images
|
||||
|
||||
## [1.0.14-alfa]
|
||||
|
||||
### Added
|
||||
- New release script added to tag images with release number
|
||||
- Allow the addition of multiple types of Catalogs
|
||||
- Generic functionality to enable dynamic fields
|
||||
- Addition of Retrievers to allow for smart collection of information in Catalogs
|
||||
- Add dynamic fields to Catalog / Retriever / DocumentVersion
|
||||
|
||||
### Changed
|
||||
- Processing parameters defined at Catalog level iso Tenant level
|
||||
- Reroute 'blank' paths to 'admin'
|
||||
|
||||
### Deprecated
|
||||
- For soon-to-be removed features.
|
||||
|
||||
### Removed
|
||||
- For now removed features.
|
||||
|
||||
### Fixed
|
||||
- Set default language when registering Documents or URLs.
|
||||
|
||||
|
||||
85
Evie Overview.md
Normal file
85
Evie Overview.md
Normal file
@@ -0,0 +1,85 @@
|
||||
# Evie Overview
|
||||
|
||||
Owner: pieter Laroy
|
||||
|
||||
# Introduction
|
||||
|
||||
The Evie project (developed by AskEveAI) is a SAAS product that enables SMEs to easily introduce AI optimisations for both internal and external use. There are two big concepts:
|
||||
|
||||
- Catalogs: these allow tenants to store information about their organisations or enterprises
|
||||
- Specialists: these allow tenants to perform logic into their processes, communications, …
|
||||
|
||||
As such, we could say we have an advanced RAG system tenants can use to optimise their workings.
|
||||
|
||||
## Multi-tenant
|
||||
|
||||
The application has a multi-tenant setup built in. This is reflected in:
|
||||
|
||||
- The Database:
|
||||
- We have 1 public schema, in which general information is defined such as tenants, their users, domains, licenses, …
|
||||
- We have a schema (named 1, 2, …) for each of the tenants defined in the system, containing all information on the tenant’s catalogs & documents, specialists & interactions, …
|
||||
- File Storage
|
||||
- We use S3-compatible storage
|
||||
- A bucket is defined for each tenant, storing their specific documents, assets, …
|
||||
|
||||
That way, general information required for the operation of Evie is stored in the public schema, and specific and potentially sensitive information is nicely stored behind a Chinese wall for each of the tenants.
|
||||
|
||||
## Partners
|
||||
|
||||
We started to define the concept of a partner. This allows us to have partners that introduce tenants to Evie, or offer them additional functionality (specialists) or knowledge (catalogs). This concept is in an early stage at this point.
|
||||
|
||||
## Domains
|
||||
|
||||
In order to ensure a structured approach, we have defined several domains in the project:
|
||||
|
||||
- **User**: the user domain is used to store all data on partners, tenants, actual users.
|
||||
- **Document**: the document domain is used to store all information on catalogs, documents, how to process documents, …
|
||||
- **Interaction**: This domain allows us to define specialists, agents, … and to interact with the specialists and agents.
|
||||
- **Entitlements**: This domain defines all license information, usage, …
|
||||
|
||||
# Project Structure
|
||||
|
||||
## Common
|
||||
|
||||
The common folder contains code that is used in different components of the system. It contains the following important pieces:
|
||||
|
||||
- **models**: in the models folder you can find the SQLAlchemy models used throughout the application. These models are organised in their relevant domains.
|
||||
- **eveai_model**: some classes to handle usage, wrappers around standard LLM clients
|
||||
- **langchain**: similar to eveai_model, but in the langchain library
|
||||
- **services**: I started to define services to define reusable functionality in the system. There again are defined in their respective domains
|
||||
- **utils**: a whole bunch of utility classes. Some should get converted to services classes in the future
|
||||
- **utils/cache**: contains code for caching different elements in the application
|
||||
|
||||
## config
|
||||
|
||||
The config folder contains quite some configuration data (as the name suggests):
|
||||
|
||||
- **config.py**: general configuration
|
||||
- **logging_config.py**: definition of logging files
|
||||
- **model_config.py**: obsolete
|
||||
- **type_defs**: contains the lists of definitions for several types used throughout the application. E.g. processor_types, specialist_types, …
|
||||
- **All other folders**: detailed configuration of all the types defined in type_defs.
|
||||
|
||||
## docker
|
||||
|
||||
The docker folder contains the configuration and scripts used for all operations on configuring and building containers, distributing containers, …
|
||||
|
||||
## eveai_… folders
|
||||
|
||||
These are different components (containerized) of our application:
|
||||
|
||||
- **eveai_api**: The API of our application.
|
||||
- **eveai_app**: The administrative interface of our application.
|
||||
- **eveai_beat**: a means to install batch processes for our application.
|
||||
- **eveai_chat**: obsolete at this moment
|
||||
- **eveai_chat_workers**: celery based invocation of our specialists
|
||||
- **eveai_client**: newly added. A desktop client to invoking specialists.
|
||||
- **eveai_entitlements**: celery based approach to handling business events, measuring and updating usage, …
|
||||
- **eveai_workers**: celery based approach to filling catalogs with documents (embedding)
|
||||
|
||||
## Remaining folders
|
||||
|
||||
- **integrations**: integrations to e.g. Wordpress and Zapier.
|
||||
- **migrations**: SQLAlchemy database migration files (for public and tenant schema)
|
||||
- **nginx**: configuration and static files for nginx
|
||||
- **scripts**: various scripts used to start up components, to perform database operations, …
|
||||
BIN
common/.DS_Store
vendored
BIN
common/.DS_Store
vendored
Binary file not shown.
Binary file not shown.
Binary file not shown.
11
common/eveai_model/eveai_embedding_base.py
Normal file
11
common/eveai_model/eveai_embedding_base.py
Normal file
@@ -0,0 +1,11 @@
|
||||
from abc import abstractmethod
|
||||
from typing import List
|
||||
|
||||
|
||||
class EveAIEmbeddings:
|
||||
@abstractmethod
|
||||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||||
raise NotImplementedError
|
||||
|
||||
def embed_query(self, text: str) -> List[float]:
|
||||
return self.embed_documents([text])[0]
|
||||
141
common/eveai_model/tracked_mistral_embeddings.py
Normal file
141
common/eveai_model/tracked_mistral_embeddings.py
Normal file
@@ -0,0 +1,141 @@
|
||||
from flask import current_app
|
||||
from langchain_mistralai import MistralAIEmbeddings
|
||||
from typing import List, Any
|
||||
import time
|
||||
|
||||
from common.eveai_model.eveai_embedding_base import EveAIEmbeddings
|
||||
from common.utils.business_event_context import current_event
|
||||
from mistralai import Mistral
|
||||
|
||||
|
||||
class TrackedMistralAIEmbeddings(EveAIEmbeddings):
|
||||
def __init__(self, model: str = "mistral_embed", batch_size: int = 10):
|
||||
"""
|
||||
Initialize the TrackedMistralAIEmbeddings class.
|
||||
|
||||
Args:
|
||||
model: The embedding model to use
|
||||
batch_size: Maximum number of texts to send in a single API call
|
||||
"""
|
||||
api_key = current_app.config['MISTRAL_API_KEY']
|
||||
self.client = Mistral(
|
||||
api_key=api_key
|
||||
)
|
||||
self.model = model
|
||||
self.batch_size = batch_size
|
||||
super().__init__()
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
"""
|
||||
Embed a list of texts, processing in batches to avoid API limitations.
|
||||
|
||||
Args:
|
||||
texts: A list of texts to embed
|
||||
|
||||
Returns:
|
||||
A list of embeddings, one for each input text
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
all_embeddings = []
|
||||
|
||||
# Process texts in batches
|
||||
for i in range(0, len(texts), self.batch_size):
|
||||
batch = texts[i:i + self.batch_size]
|
||||
batch_num = i // self.batch_size + 1
|
||||
current_app.logger.debug(f"Processing embedding batch {batch_num}, size: {len(batch)}")
|
||||
|
||||
start_time = time.time()
|
||||
try:
|
||||
result = self.client.embeddings.create(
|
||||
model=self.model,
|
||||
inputs=batch
|
||||
)
|
||||
end_time = time.time()
|
||||
batch_time = end_time - start_time
|
||||
|
||||
batch_embeddings = [embedding.embedding for embedding in result.data]
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
# Log metrics for this batch
|
||||
metrics = {
|
||||
'total_tokens': result.usage.total_tokens,
|
||||
'prompt_tokens': result.usage.prompt_tokens,
|
||||
'completion_tokens': result.usage.completion_tokens,
|
||||
'time_elapsed': batch_time,
|
||||
'interaction_type': 'Embedding',
|
||||
'batch': batch_num,
|
||||
'batch_size': len(batch)
|
||||
}
|
||||
current_event.log_llm_metrics(metrics)
|
||||
|
||||
current_app.logger.debug(f"Batch {batch_num} processed: {len(batch)} texts, "
|
||||
f"{result.usage.total_tokens} tokens, {batch_time:.2f}s")
|
||||
|
||||
# If processing multiple batches, add a small delay to avoid rate limits
|
||||
if len(texts) > self.batch_size and i + self.batch_size < len(texts):
|
||||
time.sleep(0.25) # 250ms pause between batches
|
||||
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error in embedding batch {batch_num}: {str(e)}")
|
||||
# If a batch fails, try to process each text individually
|
||||
for j, text in enumerate(batch):
|
||||
try:
|
||||
current_app.logger.debug(f"Attempting individual embedding for item {i + j}")
|
||||
single_start_time = time.time()
|
||||
single_result = self.client.embeddings.create(
|
||||
model=self.model,
|
||||
inputs=[text]
|
||||
)
|
||||
single_end_time = time.time()
|
||||
|
||||
# Add the single embedding
|
||||
single_embedding = single_result.data[0].embedding
|
||||
all_embeddings.append(single_embedding)
|
||||
|
||||
# Log metrics for this individual embedding
|
||||
single_metrics = {
|
||||
'total_tokens': single_result.usage.total_tokens,
|
||||
'prompt_tokens': single_result.usage.prompt_tokens,
|
||||
'completion_tokens': single_result.usage.completion_tokens,
|
||||
'time_elapsed': single_end_time - single_start_time,
|
||||
'interaction_type': 'Embedding',
|
||||
'batch': f"{batch_num}-recovery-{j}",
|
||||
'batch_size': 1
|
||||
}
|
||||
current_event.log_llm_metrics(single_metrics)
|
||||
|
||||
except Exception as inner_e:
|
||||
current_app.logger.error(f"Failed to embed individual text at index {i + j}: {str(inner_e)}")
|
||||
# Add a zero vector as a placeholder for failed embeddings
|
||||
# Use the correct dimensionality for the model (1024 for mistral_embed)
|
||||
embedding_dim = 1024
|
||||
all_embeddings.append([0.0] * embedding_dim)
|
||||
|
||||
total_batches = (len(texts) + self.batch_size - 1) // self.batch_size
|
||||
current_app.logger.info(f"Embedded {len(texts)} texts in {total_batches} batches")
|
||||
|
||||
return all_embeddings
|
||||
|
||||
# def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
# start_time = time.time()
|
||||
# result = self.client.embeddings.create(
|
||||
# model=self.model,
|
||||
# inputs=texts
|
||||
# )
|
||||
# end_time = time.time()
|
||||
#
|
||||
# metrics = {
|
||||
# 'total_tokens': result.usage.total_tokens,
|
||||
# 'prompt_tokens': result.usage.prompt_tokens, # For embeddings, all tokens are prompt tokens
|
||||
# 'completion_tokens': result.usage.completion_tokens,
|
||||
# 'time_elapsed': end_time - start_time,
|
||||
# 'interaction_type': 'Embedding',
|
||||
# }
|
||||
# current_event.log_llm_metrics(metrics)
|
||||
#
|
||||
# embeddings = [embedding.embedding for embedding in result.data]
|
||||
#
|
||||
# return embeddings
|
||||
|
||||
53
common/eveai_model/tracked_mistral_ocr_client.py
Normal file
53
common/eveai_model/tracked_mistral_ocr_client.py
Normal file
@@ -0,0 +1,53 @@
|
||||
import re
|
||||
import time
|
||||
|
||||
from flask import current_app
|
||||
from mistralai import Mistral
|
||||
|
||||
from common.utils.business_event_context import current_event
|
||||
|
||||
|
||||
class TrackedMistralOcrClient:
|
||||
def __init__(self):
|
||||
api_key = current_app.config['MISTRAL_API_KEY']
|
||||
self.client = Mistral(
|
||||
api_key=api_key,
|
||||
)
|
||||
self.model = "mistral-ocr-latest"
|
||||
|
||||
def _get_title(self, markdown):
|
||||
# Look for the first level-1 heading
|
||||
match = re.search(r'^# (.+)', markdown, re.MULTILINE)
|
||||
return match.group(1).strip() if match else None
|
||||
|
||||
def process_pdf(self, file_name, file_content):
|
||||
start_time = time.time()
|
||||
uploaded_pdf = self.client.files.upload(
|
||||
file={
|
||||
"file_name": file_name,
|
||||
"content": file_content
|
||||
},
|
||||
purpose="ocr"
|
||||
)
|
||||
signed_url = self.client.files.get_signed_url(file_id=uploaded_pdf.id)
|
||||
ocr_response = self.client.ocr.process(
|
||||
model=self.model,
|
||||
document={
|
||||
"type": "document_url",
|
||||
"document_url": signed_url.url
|
||||
},
|
||||
include_image_base64=False
|
||||
)
|
||||
nr_of_pages = len(ocr_response.pages)
|
||||
all_markdown = " ".join(page.markdown for page in ocr_response.pages)
|
||||
title = self._get_title(all_markdown)
|
||||
end_time = time.time()
|
||||
|
||||
metrics = {
|
||||
'nr_of_pages': nr_of_pages,
|
||||
'time_elapsed': end_time - start_time,
|
||||
'interaction_type': 'OCR',
|
||||
}
|
||||
current_event.log_llm_metrics(metrics)
|
||||
|
||||
return all_markdown, title
|
||||
@@ -2,16 +2,15 @@ from flask_sqlalchemy import SQLAlchemy
|
||||
from flask_migrate import Migrate
|
||||
from flask_bootstrap import Bootstrap
|
||||
from flask_security import Security
|
||||
from flask_mailman import Mail
|
||||
from flask_login import LoginManager
|
||||
from flask_cors import CORS
|
||||
from flask_socketio import SocketIO
|
||||
from flask_jwt_extended import JWTManager
|
||||
from flask_session import Session
|
||||
from flask_wtf import CSRFProtect
|
||||
from flask_restx import Api
|
||||
from prometheus_flask_exporter import PrometheusMetrics
|
||||
|
||||
from .utils.cache.eveai_cache_manager import EveAICacheManager
|
||||
from .utils.simple_encryption import SimpleEncryption
|
||||
from .utils.minio_utils import MinioClient
|
||||
|
||||
@@ -22,13 +21,13 @@ migrate = Migrate()
|
||||
bootstrap = Bootstrap()
|
||||
csrf = CSRFProtect()
|
||||
security = Security()
|
||||
mail = Mail()
|
||||
login_manager = LoginManager()
|
||||
cors = CORS()
|
||||
socketio = SocketIO()
|
||||
jwt = JWTManager()
|
||||
session = Session()
|
||||
api_rest = Api()
|
||||
simple_encryption = SimpleEncryption()
|
||||
minio_client = MinioClient()
|
||||
metrics = PrometheusMetrics.for_app_factory()
|
||||
cache_manager = EveAICacheManager()
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
23
common/langchain/outputs/base.py
Normal file
23
common/langchain/outputs/base.py
Normal file
@@ -0,0 +1,23 @@
|
||||
# Output Schema Management - common/langchain/outputs/base.py
|
||||
from typing import Dict, Type, Any
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class BaseSpecialistOutput(BaseModel):
|
||||
"""Base class for all specialist outputs"""
|
||||
pass
|
||||
|
||||
|
||||
class OutputRegistry:
|
||||
"""Registry for specialist output schemas"""
|
||||
_schemas: Dict[str, Type[BaseSpecialistOutput]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, specialist_type: str, schema_class: Type[BaseSpecialistOutput]):
|
||||
cls._schemas[specialist_type] = schema_class
|
||||
|
||||
@classmethod
|
||||
def get_schema(cls, specialist_type: str) -> Type[BaseSpecialistOutput]:
|
||||
if specialist_type not in cls._schemas:
|
||||
raise ValueError(f"No output schema registered for {specialist_type}")
|
||||
return cls._schemas[specialist_type]
|
||||
22
common/langchain/outputs/rag.py
Normal file
22
common/langchain/outputs/rag.py
Normal file
@@ -0,0 +1,22 @@
|
||||
# RAG Specialist Output - common/langchain/outputs/rag.py
|
||||
from typing import List
|
||||
from pydantic import Field
|
||||
from .base import BaseSpecialistOutput
|
||||
|
||||
|
||||
class RAGOutput(BaseSpecialistOutput):
|
||||
"""Output schema for RAG specialist"""
|
||||
"""Default docstring - to be replaced with actual prompt"""
|
||||
|
||||
answer: str = Field(
|
||||
...,
|
||||
description="The answer to the user question, based on the given sources",
|
||||
)
|
||||
citations: List[int] = Field(
|
||||
...,
|
||||
description="The integer IDs of the SPECIFIC sources that were used to generate the answer"
|
||||
)
|
||||
insufficient_info: bool = Field(
|
||||
False, # Default value is set to False
|
||||
description="A boolean indicating whether given sources were sufficient or not to generate the answer"
|
||||
)
|
||||
@@ -1,145 +0,0 @@
|
||||
from langchain_core.retrievers import BaseRetriever
|
||||
from sqlalchemy import func, and_, or_, desc
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from pydantic import BaseModel, Field, PrivateAttr
|
||||
from typing import Any, Dict
|
||||
from flask import current_app
|
||||
|
||||
from common.extensions import db
|
||||
from common.models.document import Document, DocumentVersion
|
||||
from common.utils.datetime_utils import get_date_in_timezone
|
||||
from common.utils.model_utils import ModelVariables
|
||||
|
||||
|
||||
class EveAIDefaultRagRetriever(BaseRetriever, BaseModel):
|
||||
_catalog_id: int = PrivateAttr()
|
||||
_model_variables: ModelVariables = PrivateAttr()
|
||||
_tenant_info: Dict[str, Any] = PrivateAttr()
|
||||
|
||||
def __init__(self, catalog_id: int, model_variables: ModelVariables, tenant_info: Dict[str, Any]):
|
||||
super().__init__()
|
||||
current_app.logger.debug(f'Model variables type: {type(model_variables)}')
|
||||
self._catalog_id = catalog_id
|
||||
self._model_variables = model_variables
|
||||
self._tenant_info = tenant_info
|
||||
|
||||
@property
|
||||
def catalog_id(self) -> int:
|
||||
return self._catalog_id
|
||||
|
||||
@property
|
||||
def model_variables(self) -> ModelVariables:
|
||||
return self._model_variables
|
||||
|
||||
@property
|
||||
def tenant_info(self) -> Dict[str, Any]:
|
||||
return self._tenant_info
|
||||
|
||||
def _get_relevant_documents(self, query: str):
|
||||
current_app.logger.debug(f'Retrieving relevant documents for query: {query}')
|
||||
query_embedding = self._get_query_embedding(query)
|
||||
current_app.logger.debug(f'Model Variables Private: {type(self._model_variables)}')
|
||||
current_app.logger.debug(f'Model Variables Property: {type(self.model_variables)}')
|
||||
db_class = self.model_variables['embedding_db_model']
|
||||
similarity_threshold = self.model_variables['similarity_threshold']
|
||||
k = self.model_variables['k']
|
||||
|
||||
if self.model_variables['rag_tuning']:
|
||||
try:
|
||||
current_date = get_date_in_timezone(self.tenant_info['timezone'])
|
||||
current_app.rag_tuning_logger.debug(f'Current date: {current_date}\n')
|
||||
|
||||
# Debug query to show similarity for all valid documents (without chunk text)
|
||||
debug_query = (
|
||||
db.session.query(
|
||||
Document.id.label('document_id'),
|
||||
DocumentVersion.id.label('version_id'),
|
||||
db_class.id.label('embedding_id'),
|
||||
(1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity')
|
||||
)
|
||||
.join(DocumentVersion, db_class.doc_vers_id == DocumentVersion.id)
|
||||
.join(Document, DocumentVersion.doc_id == Document.id)
|
||||
.filter(
|
||||
or_(Document.valid_from.is_(None), func.date(Document.valid_from) <= current_date),
|
||||
or_(Document.valid_to.is_(None), func.date(Document.valid_to) >= current_date)
|
||||
)
|
||||
.order_by(desc('similarity'))
|
||||
)
|
||||
|
||||
debug_results = debug_query.all()
|
||||
|
||||
current_app.logger.debug("Debug: Similarity for all valid documents:")
|
||||
for row in debug_results:
|
||||
current_app.rag_tuning_logger.debug(f"Doc ID: {row.document_id}, "
|
||||
f"Version ID: {row.version_id}, "
|
||||
f"Embedding ID: {row.embedding_id}, "
|
||||
f"Similarity: {row.similarity}")
|
||||
current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f'Error generating overview: {e}')
|
||||
db.session.rollback()
|
||||
|
||||
if self.model_variables['rag_tuning']:
|
||||
current_app.rag_tuning_logger.debug(f'Parameters for Retrieval of documents: \n')
|
||||
current_app.rag_tuning_logger.debug(f'Similarity Threshold: {similarity_threshold}\n')
|
||||
current_app.rag_tuning_logger.debug(f'K: {k}\n')
|
||||
current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
|
||||
|
||||
try:
|
||||
current_date = get_date_in_timezone(self.tenant_info['timezone'])
|
||||
# Subquery to find the latest version of each document
|
||||
subquery = (
|
||||
db.session.query(
|
||||
DocumentVersion.doc_id,
|
||||
func.max(DocumentVersion.id).label('latest_version_id')
|
||||
)
|
||||
.group_by(DocumentVersion.doc_id)
|
||||
.subquery()
|
||||
)
|
||||
# Main query to filter embeddings
|
||||
query_obj = (
|
||||
db.session.query(db_class,
|
||||
(1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity'))
|
||||
.join(DocumentVersion, db_class.doc_vers_id == DocumentVersion.id)
|
||||
.join(Document, DocumentVersion.doc_id == Document.id)
|
||||
.join(subquery, DocumentVersion.id == subquery.c.latest_version_id)
|
||||
.filter(
|
||||
or_(Document.valid_from.is_(None), func.date(Document.valid_from) <= current_date),
|
||||
or_(Document.valid_to.is_(None), func.date(Document.valid_to) >= current_date),
|
||||
(1 - db_class.embedding.cosine_distance(query_embedding)) > similarity_threshold,
|
||||
Document.catalog_id == self._catalog_id
|
||||
)
|
||||
.order_by(desc('similarity'))
|
||||
.limit(k)
|
||||
)
|
||||
|
||||
if self.model_variables['rag_tuning']:
|
||||
current_app.rag_tuning_logger.debug(f'Query executed for Retrieval of documents: \n')
|
||||
current_app.rag_tuning_logger.debug(f'{query_obj.statement}\n')
|
||||
current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
|
||||
|
||||
res = query_obj.all()
|
||||
|
||||
if self.model_variables['rag_tuning']:
|
||||
current_app.rag_tuning_logger.debug(f'Retrieved {len(res)} relevant documents \n')
|
||||
current_app.rag_tuning_logger.debug(f'Data retrieved: \n')
|
||||
current_app.rag_tuning_logger.debug(f'{res}\n')
|
||||
current_app.rag_tuning_logger.debug(f'---------------------------------------\n')
|
||||
|
||||
result = []
|
||||
for doc in res:
|
||||
if self.model_variables['rag_tuning']:
|
||||
current_app.rag_tuning_logger.debug(f'Document ID: {doc[0].id} - Distance: {doc[1]}\n')
|
||||
current_app.rag_tuning_logger.debug(f'Chunk: \n {doc[0].chunk}\n\n')
|
||||
result.append(f'SOURCE: {doc[0].id}\n\n{doc[0].chunk}\n\n')
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f'Error retrieving relevant documents: {e}')
|
||||
db.session.rollback()
|
||||
return []
|
||||
return result
|
||||
|
||||
def _get_query_embedding(self, query: str):
|
||||
embedding_model = self.model_variables['embedding_model']
|
||||
query_embedding = embedding_model.embed_query(query)
|
||||
return query_embedding
|
||||
@@ -1,154 +0,0 @@
|
||||
from langchain_core.retrievers import BaseRetriever
|
||||
from sqlalchemy import func, and_, or_, desc, cast, JSON
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from pydantic import BaseModel, Field, PrivateAttr
|
||||
from typing import Any, Dict, List, Optional
|
||||
from flask import current_app
|
||||
from contextlib import contextmanager
|
||||
|
||||
from common.extensions import db
|
||||
from common.models.document import Document, DocumentVersion, Catalog
|
||||
from common.utils.datetime_utils import get_date_in_timezone
|
||||
from common.utils.model_utils import ModelVariables
|
||||
|
||||
|
||||
class EveAIDossierRetriever(BaseRetriever, BaseModel):
|
||||
_catalog_id: int = PrivateAttr()
|
||||
_model_variables: ModelVariables = PrivateAttr()
|
||||
_tenant_info: Dict[str, Any] = PrivateAttr()
|
||||
_active_filters: Optional[Dict[str, Any]] = PrivateAttr()
|
||||
|
||||
def __init__(self, catalog_id: int, model_variables: ModelVariables, tenant_info: Dict[str, Any]):
|
||||
super().__init__()
|
||||
self._catalog_id = catalog_id
|
||||
self._model_variables = model_variables
|
||||
self._tenant_info = tenant_info
|
||||
self._active_filters = None
|
||||
|
||||
@contextmanager
|
||||
def filtering(self, metadata_filters: Dict[str, Any]):
|
||||
"""Context manager for temporarily setting metadata filters"""
|
||||
previous_filters = self._active_filters
|
||||
self._active_filters = metadata_filters
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._active_filters = previous_filters
|
||||
|
||||
def _build_metadata_filter_conditions(self, query):
|
||||
"""Build SQL conditions for metadata filtering"""
|
||||
if not self._active_filters:
|
||||
return query
|
||||
|
||||
conditions = []
|
||||
for field, value in self._active_filters.items():
|
||||
if value is None:
|
||||
continue
|
||||
|
||||
# Handle both single values and lists of values
|
||||
if isinstance(value, (list, tuple)):
|
||||
# Multiple values - create OR condition
|
||||
or_conditions = []
|
||||
for val in value:
|
||||
or_conditions.append(
|
||||
cast(DocumentVersion.user_metadata[field].astext, JSON) == str(val)
|
||||
)
|
||||
if or_conditions:
|
||||
conditions.append(or_(*or_conditions))
|
||||
else:
|
||||
# Single value - direct comparison
|
||||
conditions.append(
|
||||
cast(DocumentVersion.user_metadata[field].astext, JSON) == str(value)
|
||||
)
|
||||
|
||||
if conditions:
|
||||
query = query.filter(and_(*conditions))
|
||||
|
||||
return query
|
||||
|
||||
def _get_relevant_documents(self, query: str):
|
||||
current_app.logger.debug(f'Retrieving relevant documents for dossier query: {query}')
|
||||
if self._active_filters:
|
||||
current_app.logger.debug(f'Using metadata filters: {self._active_filters}')
|
||||
|
||||
query_embedding = self._get_query_embedding(query)
|
||||
db_class = self.model_variables['embedding_db_model']
|
||||
similarity_threshold = self.model_variables['similarity_threshold']
|
||||
k = self.model_variables['k']
|
||||
|
||||
try:
|
||||
current_date = get_date_in_timezone(self.tenant_info['timezone'])
|
||||
|
||||
# Subquery to find the latest version of each document
|
||||
subquery = (
|
||||
db.session.query(
|
||||
DocumentVersion.doc_id,
|
||||
func.max(DocumentVersion.id).label('latest_version_id')
|
||||
)
|
||||
.group_by(DocumentVersion.doc_id)
|
||||
.subquery()
|
||||
)
|
||||
|
||||
# Build base query
|
||||
# Build base query
|
||||
query_obj = (
|
||||
db.session.query(db_class,
|
||||
(1 - db_class.embedding.cosine_distance(query_embedding)).label('similarity'))
|
||||
.join(DocumentVersion, db_class.doc_vers_id == DocumentVersion.id)
|
||||
.join(Document, DocumentVersion.doc_id == Document.id)
|
||||
.join(subquery, DocumentVersion.id == subquery.c.latest_version_id)
|
||||
.filter(
|
||||
or_(Document.valid_from.is_(None), func.date(Document.valid_from) <= current_date),
|
||||
or_(Document.valid_to.is_(None), func.date(Document.valid_to) >= current_date),
|
||||
(1 - db_class.embedding.cosine_distance(query_embedding)) > similarity_threshold,
|
||||
Document.catalog_id == self._catalog_id
|
||||
)
|
||||
)
|
||||
|
||||
# Apply metadata filters
|
||||
query_obj = self._build_metadata_filter_conditions(query_obj)
|
||||
|
||||
# Order and limit results
|
||||
query_obj = query_obj.order_by(desc('similarity')).limit(k)
|
||||
|
||||
# Debug logging for RAG tuning if enabled
|
||||
if self.model_variables['rag_tuning']:
|
||||
self._log_rag_tuning(query_obj, query_embedding)
|
||||
|
||||
res = query_obj.all()
|
||||
|
||||
result = []
|
||||
for doc in res:
|
||||
if self.model_variables['rag_tuning']:
|
||||
current_app.logger.debug(f'Document ID: {doc[0].id} - Distance: {doc[1]}\n')
|
||||
current_app.logger.debug(f'Chunk: \n {doc[0].chunk}\n\n')
|
||||
result.append(f'SOURCE: {doc[0].id}\n\n{doc[0].chunk}\n\n')
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f'Error retrieving relevant documents: {e}')
|
||||
db.session.rollback()
|
||||
return []
|
||||
|
||||
return result
|
||||
|
||||
def _log_rag_tuning(self, query_obj, query_embedding):
|
||||
"""Log debug information for RAG tuning"""
|
||||
current_app.rag_tuning_logger.debug("Debug: Query execution plan:")
|
||||
current_app.rag_tuning_logger.debug(f"{query_obj.statement}")
|
||||
if self._active_filters:
|
||||
current_app.rag_tuning_logger.debug("Debug: Active metadata filters:")
|
||||
current_app.rag_tuning_logger.debug(f"{self._active_filters}")
|
||||
|
||||
def _get_query_embedding(self, query: str):
|
||||
"""Get embedding for the query text"""
|
||||
embedding_model = self.model_variables['embedding_model']
|
||||
query_embedding = embedding_model.embed_query(query)
|
||||
return query_embedding
|
||||
|
||||
@property
|
||||
def model_variables(self) -> ModelVariables:
|
||||
return self._model_variables
|
||||
|
||||
@property
|
||||
def tenant_info(self) -> Dict[str, Any]:
|
||||
return self._tenant_info
|
||||
@@ -1,52 +0,0 @@
|
||||
from langchain_core.retrievers import BaseRetriever
|
||||
from sqlalchemy import asc
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from pydantic import Field, BaseModel, PrivateAttr
|
||||
from typing import Any, Dict
|
||||
from flask import current_app
|
||||
|
||||
from common.extensions import db
|
||||
from common.models.interaction import ChatSession, Interaction
|
||||
from common.utils.model_utils import ModelVariables
|
||||
|
||||
|
||||
class EveAIHistoryRetriever(BaseRetriever, BaseModel):
|
||||
_model_variables: ModelVariables = PrivateAttr()
|
||||
_session_id: str = PrivateAttr()
|
||||
|
||||
def __init__(self, model_variables: ModelVariables, session_id: str):
|
||||
super().__init__()
|
||||
self._model_variables = model_variables
|
||||
self._session_id = session_id
|
||||
|
||||
@property
|
||||
def model_variables(self) -> ModelVariables:
|
||||
return self._model_variables
|
||||
|
||||
@property
|
||||
def session_id(self) -> str:
|
||||
return self._session_id
|
||||
|
||||
def _get_relevant_documents(self, query: str):
|
||||
current_app.logger.debug(f'Retrieving history of interactions for query: {query}')
|
||||
|
||||
try:
|
||||
query_obj = (
|
||||
db.session.query(Interaction)
|
||||
.join(ChatSession, Interaction.chat_session_id == ChatSession.id)
|
||||
.filter(ChatSession.session_id == self.session_id)
|
||||
.order_by(asc(Interaction.id))
|
||||
)
|
||||
|
||||
interactions = query_obj.all()
|
||||
|
||||
result = []
|
||||
for interaction in interactions:
|
||||
result.append(f'HUMAN:\n{interaction.detailed_question}\n\nAI: \n{interaction.answer}\n\n')
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f'Error retrieving history of interactions: {e}')
|
||||
db.session.rollback()
|
||||
return []
|
||||
|
||||
return result
|
||||
@@ -1,40 +0,0 @@
|
||||
from pydantic import BaseModel, PrivateAttr
|
||||
from typing import Dict, Any
|
||||
|
||||
from common.utils.model_utils import ModelVariables
|
||||
|
||||
|
||||
class EveAIRetriever(BaseModel):
|
||||
_catalog_id: int = PrivateAttr()
|
||||
_user_metadata: Dict[str, Any] = PrivateAttr()
|
||||
_system_metadata: Dict[str, Any] = PrivateAttr()
|
||||
_configuration: Dict[str, Any] = PrivateAttr()
|
||||
_tenant_info: Dict[str, Any] = PrivateAttr()
|
||||
_model_variables: ModelVariables = PrivateAttr()
|
||||
_tuning: bool = PrivateAttr()
|
||||
|
||||
def __init__(self, catalog_id: int, user_metadata: Dict[str, Any], system_metadata: Dict[str, Any],
|
||||
configuration: Dict[str, Any]):
|
||||
super().__init__()
|
||||
self._catalog_id = catalog_id
|
||||
self._user_metadata = user_metadata
|
||||
self._system_metadata = system_metadata
|
||||
self._configuration = configuration
|
||||
|
||||
@property
|
||||
def catalog_id(self):
|
||||
return self._catalog_id
|
||||
|
||||
@property
|
||||
def user_metadata(self):
|
||||
return self._user_metadata
|
||||
|
||||
@property
|
||||
def system_metadata(self):
|
||||
return self._system_metadata
|
||||
|
||||
@property
|
||||
def configuration(self):
|
||||
return self._configuration
|
||||
|
||||
# Any common methods that should be shared among retrievers can go here.
|
||||
@@ -1,27 +0,0 @@
|
||||
import time
|
||||
from common.utils.business_event_context import current_event
|
||||
|
||||
|
||||
def tracked_transcribe(client, *args, **kwargs):
|
||||
start_time = time.time()
|
||||
|
||||
# Extract the file and model from kwargs if present, otherwise use defaults
|
||||
file = kwargs.get('file')
|
||||
model = kwargs.get('model', 'whisper-1')
|
||||
duration = kwargs.pop('duration', 600)
|
||||
|
||||
result = client.audio.transcriptions.create(*args, **kwargs)
|
||||
end_time = time.time()
|
||||
|
||||
# Token usage for transcriptions is actually the duration in seconds we pass, as the whisper model is priced per second transcribed
|
||||
|
||||
metrics = {
|
||||
'total_tokens': duration,
|
||||
'prompt_tokens': 0, # For transcriptions, all tokens are considered "completion"
|
||||
'completion_tokens': duration,
|
||||
'time_elapsed': end_time - start_time,
|
||||
'interaction_type': 'ASR',
|
||||
}
|
||||
current_event.log_llm_metrics(metrics)
|
||||
|
||||
return result
|
||||
77
common/langchain/tracked_transcription.py
Normal file
77
common/langchain/tracked_transcription.py
Normal file
@@ -0,0 +1,77 @@
|
||||
# common/langchain/tracked_transcription.py
|
||||
from typing import Any, Optional, Dict
|
||||
import time
|
||||
from openai import OpenAI
|
||||
from common.utils.business_event_context import current_event
|
||||
|
||||
|
||||
class TrackedOpenAITranscription:
|
||||
"""Wrapper for OpenAI transcription with metric tracking"""
|
||||
|
||||
def __init__(self, api_key: str, **kwargs: Any):
|
||||
"""Initialize with OpenAI client settings"""
|
||||
self.client = OpenAI(api_key=api_key)
|
||||
self.model = kwargs.get('model', 'whisper-1')
|
||||
|
||||
def transcribe(self,
|
||||
file: Any,
|
||||
model: Optional[str] = None,
|
||||
language: Optional[str] = None,
|
||||
prompt: Optional[str] = None,
|
||||
response_format: Optional[str] = None,
|
||||
temperature: Optional[float] = None,
|
||||
duration: Optional[int] = None) -> str:
|
||||
"""
|
||||
Transcribe audio with metrics tracking
|
||||
|
||||
Args:
|
||||
file: Audio file to transcribe
|
||||
model: Model to use (defaults to whisper-1)
|
||||
language: Optional language of the audio
|
||||
prompt: Optional prompt to guide transcription
|
||||
response_format: Response format (json, text, etc)
|
||||
temperature: Sampling temperature
|
||||
duration: Duration of audio in seconds for metrics
|
||||
|
||||
Returns:
|
||||
Transcription text
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Create transcription options
|
||||
options = {
|
||||
"file": file,
|
||||
"model": model or self.model,
|
||||
}
|
||||
if language:
|
||||
options["language"] = language
|
||||
if prompt:
|
||||
options["prompt"] = prompt
|
||||
if response_format:
|
||||
options["response_format"] = response_format
|
||||
if temperature:
|
||||
options["temperature"] = temperature
|
||||
|
||||
response = self.client.audio.transcriptions.create(**options)
|
||||
|
||||
# Calculate metrics
|
||||
end_time = time.time()
|
||||
|
||||
# Token usage for transcriptions is based on audio duration
|
||||
metrics = {
|
||||
'total_tokens': duration or 600, # Default to 10 minutes if duration not provided
|
||||
'prompt_tokens': 0, # For transcriptions, all tokens are completion
|
||||
'completion_tokens': duration or 600,
|
||||
'time_elapsed': end_time - start_time,
|
||||
'interaction_type': 'ASR',
|
||||
}
|
||||
current_event.log_llm_metrics(metrics)
|
||||
|
||||
# Return text from response
|
||||
if isinstance(response, str):
|
||||
return response
|
||||
return response.text
|
||||
|
||||
except Exception as e:
|
||||
raise Exception(f"Transcription failed: {str(e)}")
|
||||
BIN
common/models/.DS_Store
vendored
BIN
common/models/.DS_Store
vendored
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -10,24 +10,33 @@ class Catalog(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="DEFAULT_CATALOG")
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_CATALOG")
|
||||
|
||||
# Embedding variables
|
||||
html_tags = db.Column(ARRAY(sa.String(10)), nullable=True, default=['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'li'])
|
||||
html_end_tags = db.Column(ARRAY(sa.String(10)), nullable=True, default=['p', 'li'])
|
||||
html_included_elements = db.Column(ARRAY(sa.String(50)), nullable=True)
|
||||
html_excluded_elements = db.Column(ARRAY(sa.String(50)), nullable=True)
|
||||
html_excluded_classes = db.Column(ARRAY(sa.String(200)), nullable=True)
|
||||
min_chunk_size = db.Column(db.Integer, nullable=True, default=1500)
|
||||
max_chunk_size = db.Column(db.Integer, nullable=True, default=2500)
|
||||
|
||||
min_chunk_size = db.Column(db.Integer, nullable=True, default=2000)
|
||||
max_chunk_size = db.Column(db.Integer, nullable=True, default=3000)
|
||||
# Meta Data
|
||||
user_metadata = db.Column(JSONB, nullable=True)
|
||||
system_metadata = db.Column(JSONB, nullable=True)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Chat variables ==> Move to Specialist?
|
||||
chat_RAG_temperature = db.Column(db.Float, nullable=True, default=0.3)
|
||||
chat_no_RAG_temperature = db.Column(db.Float, nullable=True, default=0.5)
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
|
||||
class Processor(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
catalog_id = db.Column(db.Integer, db.ForeignKey('catalog.id'), nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False)
|
||||
sub_file_type = db.Column(db.String(50), nullable=True)
|
||||
|
||||
# Tuning enablers
|
||||
embed_tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Meta Data
|
||||
user_metadata = db.Column(JSONB, nullable=True)
|
||||
@@ -46,13 +55,15 @@ class Retriever(db.Model):
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
catalog_id = db.Column(db.Integer, db.ForeignKey('catalog.id'), nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="DEFAULT_RAG")
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_RAG")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="STANDARD_RAG")
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Meta Data
|
||||
user_metadata = db.Column(JSONB, nullable=True)
|
||||
system_metadata = db.Column(JSONB, nullable=True)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
@@ -89,6 +100,7 @@ class DocumentVersion(db.Model):
|
||||
bucket_name = db.Column(db.String(255), nullable=True)
|
||||
object_name = db.Column(db.String(200), nullable=True)
|
||||
file_type = db.Column(db.String(20), nullable=True)
|
||||
sub_file_type = db.Column(db.String(50), nullable=True)
|
||||
file_size = db.Column(db.Float, nullable=True)
|
||||
language = db.Column(db.String(2), nullable=False)
|
||||
user_context = db.Column(db.Text, nullable=True)
|
||||
|
||||
@@ -1,4 +1,14 @@
|
||||
from sqlalchemy.sql.expression import text
|
||||
|
||||
from common.extensions import db
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
from enum import Enum
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.ext.hybrid import hybrid_property
|
||||
from dateutil.relativedelta import relativedelta
|
||||
|
||||
from common.utils.database import Database
|
||||
|
||||
|
||||
class BusinessEventLog(db.Model):
|
||||
@@ -11,10 +21,13 @@ class BusinessEventLog(db.Model):
|
||||
tenant_id = db.Column(db.Integer, nullable=False)
|
||||
trace_id = db.Column(db.String(50), nullable=False)
|
||||
span_id = db.Column(db.String(50))
|
||||
span_name = db.Column(db.String(50))
|
||||
span_name = db.Column(db.String(255))
|
||||
parent_span_id = db.Column(db.String(50))
|
||||
document_version_id = db.Column(db.Integer)
|
||||
document_version_file_size = db.Column(db.Float)
|
||||
specialist_id = db.Column(db.Integer)
|
||||
specialist_type = db.Column(db.String(50))
|
||||
specialist_type_version = db.Column(db.String(20))
|
||||
chat_session_id = db.Column(db.String(50))
|
||||
interaction_id = db.Column(db.Integer)
|
||||
environment = db.Column(db.String(20))
|
||||
@@ -22,6 +35,7 @@ class BusinessEventLog(db.Model):
|
||||
llm_metrics_prompt_tokens = db.Column(db.Integer)
|
||||
llm_metrics_completion_tokens = db.Column(db.Integer)
|
||||
llm_metrics_total_time = db.Column(db.Float)
|
||||
llm_metrics_nr_of_pages = db.Column(db.Integer)
|
||||
llm_metrics_call_count = db.Column(db.Integer)
|
||||
llm_interaction_type = db.Column(db.String(20))
|
||||
message = db.Column(db.Text)
|
||||
@@ -38,6 +52,7 @@ class License(db.Model):
|
||||
tier_id = db.Column(db.Integer, db.ForeignKey('public.license_tier.id'),nullable=False) # 'small', 'medium', 'custom'
|
||||
start_date = db.Column(db.Date, nullable=False)
|
||||
end_date = db.Column(db.Date, nullable=True)
|
||||
nr_of_periods = db.Column(db.Integer, nullable=False)
|
||||
currency = db.Column(db.String(20), nullable=False)
|
||||
yearly_payment = db.Column(db.Boolean, nullable=False, default=False)
|
||||
basic_fee = db.Column(db.Float, nullable=False)
|
||||
@@ -52,10 +67,41 @@ class License(db.Model):
|
||||
additional_interaction_bucket = db.Column(db.Integer, nullable=False)
|
||||
overage_embedding = db.Column(db.Float, nullable=False, default=0)
|
||||
overage_interaction = db.Column(db.Float, nullable=False, default=0)
|
||||
additional_storage_allowed = db.Column(db.Boolean, nullable=False, default=False)
|
||||
additional_embedding_allowed = db.Column(db.Boolean, nullable=False, default=False)
|
||||
additional_interaction_allowed = db.Column(db.Boolean, nullable=False, default=False)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
tenant = db.relationship('Tenant', back_populates='licenses')
|
||||
license_tier = db.relationship('LicenseTier', back_populates='licenses')
|
||||
usages = db.relationship('LicenseUsage', order_by='LicenseUsage.period_start_date', back_populates='license')
|
||||
periods = db.relationship('LicensePeriod', back_populates='license',
|
||||
order_by='LicensePeriod.period_number',
|
||||
cascade='all, delete-orphan')
|
||||
|
||||
def calculate_end_date(start_date, nr_of_periods):
|
||||
"""Utility functie om einddatum te berekenen"""
|
||||
if start_date and nr_of_periods:
|
||||
return start_date + relativedelta(months=nr_of_periods) - relativedelta(days=1)
|
||||
return None
|
||||
|
||||
# Luister naar start_date wijzigingen
|
||||
@event.listens_for(License.start_date, 'set')
|
||||
def set_start_date(target, value, oldvalue, initiator):
|
||||
"""Bijwerken van end_date wanneer start_date wordt aangepast"""
|
||||
if value and target.nr_of_periods:
|
||||
target.end_date = calculate_end_date(value, target.nr_of_periods)
|
||||
|
||||
# Luister naar nr_of_periods wijzigingen
|
||||
@event.listens_for(License.nr_of_periods, 'set')
|
||||
def set_nr_of_periods(target, value, oldvalue, initiator):
|
||||
"""Bijwerken van end_date wanneer nr_of_periods wordt aangepast"""
|
||||
if value and target.start_date:
|
||||
target.end_date = calculate_end_date(target.start_date, value)
|
||||
|
||||
|
||||
class LicenseTier(db.Model):
|
||||
@@ -84,7 +130,219 @@ class LicenseTier(db.Model):
|
||||
standard_overage_embedding = db.Column(db.Float, nullable=False, default=0)
|
||||
standard_overage_interaction = db.Column(db.Float, nullable=False, default=0)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
licenses = db.relationship('License', back_populates='license_tier')
|
||||
partner_services = db.relationship('PartnerServiceLicenseTier', back_populates='license_tier')
|
||||
|
||||
|
||||
class PartnerServiceLicenseTier(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
partner_service_id = db.Column(db.Integer, db.ForeignKey('public.partner_service.id'), primary_key=True,
|
||||
nullable=False)
|
||||
license_tier_id = db.Column(db.Integer, db.ForeignKey('public.license_tier.id'), primary_key=True,
|
||||
nullable=False)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
license_tier = db.relationship('LicenseTier', back_populates='partner_services')
|
||||
partner_service = db.relationship('PartnerService', back_populates='license_tiers')
|
||||
|
||||
|
||||
class PeriodStatus(Enum):
|
||||
UPCOMING = "UPCOMING" # The period is still in the future
|
||||
PENDING = "PENDING" # The period is active, but prepaid is not yet received
|
||||
ACTIVE = "ACTIVE" # The period is active and prepaid has been received
|
||||
COMPLETED = "COMPLETED" # The period has been completed, but not yet invoiced
|
||||
INVOICED = "INVOICED" # The period has been completed and invoiced, but overage payment still pending
|
||||
CLOSED = "CLOSED" # The period has been closed, invoiced and fully paid
|
||||
|
||||
|
||||
class LicensePeriod(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
license_id = db.Column(db.Integer, db.ForeignKey('public.license.id'), nullable=False)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
|
||||
# Period identification
|
||||
period_number = db.Column(db.Integer, nullable=False)
|
||||
period_start = db.Column(db.Date, nullable=False)
|
||||
period_end = db.Column(db.Date, nullable=False)
|
||||
|
||||
# License configuration snapshot - copied from license when period is created
|
||||
currency = db.Column(db.String(20), nullable=True)
|
||||
basic_fee = db.Column(db.Float, nullable=True)
|
||||
max_storage_mb = db.Column(db.Integer, nullable=True)
|
||||
additional_storage_price = db.Column(db.Float, nullable=True)
|
||||
additional_storage_bucket = db.Column(db.Integer, nullable=True)
|
||||
included_embedding_mb = db.Column(db.Integer, nullable=True)
|
||||
additional_embedding_price = db.Column(db.Numeric(10, 4), nullable=True)
|
||||
additional_embedding_bucket = db.Column(db.Integer, nullable=True)
|
||||
included_interaction_tokens = db.Column(db.Integer, nullable=True)
|
||||
additional_interaction_token_price = db.Column(db.Numeric(10, 4), nullable=True)
|
||||
additional_interaction_bucket = db.Column(db.Integer, nullable=True)
|
||||
|
||||
# Allowance flags - can be changed from False to True within a period
|
||||
additional_storage_allowed = db.Column(db.Boolean, nullable=True, default=False)
|
||||
additional_embedding_allowed = db.Column(db.Boolean, nullable=True, default=False)
|
||||
additional_interaction_allowed = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Status tracking
|
||||
status = db.Column(db.Enum(PeriodStatus), nullable=False, default=PeriodStatus.UPCOMING)
|
||||
|
||||
# State transition timestamps
|
||||
upcoming_at = db.Column(db.DateTime, nullable=True)
|
||||
pending_at = db.Column(db.DateTime, nullable=True)
|
||||
active_at = db.Column(db.DateTime, nullable=True)
|
||||
completed_at = db.Column(db.DateTime, nullable=True)
|
||||
invoiced_at = db.Column(db.DateTime, nullable=True)
|
||||
closed_at = db.Column(db.DateTime, nullable=True)
|
||||
|
||||
# Standard audit fields
|
||||
created_at = db.Column(db.DateTime, server_default=db.func.now())
|
||||
updated_at = db.Column(db.DateTime, server_default=db.func.now(), onupdate=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
|
||||
# Relationships
|
||||
license = db.relationship('License', back_populates='periods')
|
||||
license_usage = db.relationship('LicenseUsage',
|
||||
uselist=False, # This makes it one-to-one
|
||||
back_populates='license_period',
|
||||
cascade='all, delete-orphan')
|
||||
payments = db.relationship('Payment', back_populates='license_period')
|
||||
invoices = db.relationship('Invoice', back_populates='license_period',
|
||||
cascade='all, delete-orphan')
|
||||
|
||||
def update_allowance(self, allowance_type, allow_value, user_id=None):
|
||||
"""
|
||||
Update an allowance flag within a period
|
||||
Only allows transitioning from False to True
|
||||
|
||||
Args:
|
||||
allowance_type: One of 'storage', 'embedding', or 'interaction'
|
||||
allow_value: The new value (must be True)
|
||||
user_id: User ID performing the update
|
||||
|
||||
Raises:
|
||||
ValueError: If trying to change from True to False, or invalid allowance type
|
||||
"""
|
||||
field_name = f"additional_{allowance_type}_allowed"
|
||||
|
||||
# Verify valid field
|
||||
if not hasattr(self, field_name):
|
||||
raise ValueError(f"Invalid allowance type: {allowance_type}")
|
||||
|
||||
# Get current value
|
||||
current_value = getattr(self, field_name)
|
||||
|
||||
# Only allow False -> True transition
|
||||
if current_value is True and allow_value is True:
|
||||
# Already True, no change needed
|
||||
return
|
||||
elif allow_value is False:
|
||||
raise ValueError(f"Cannot change {field_name} from {current_value} to False")
|
||||
|
||||
# Update the field
|
||||
setattr(self, field_name, True)
|
||||
self.updated_at = dt.now(tz.utc)
|
||||
if user_id:
|
||||
self.updated_by = user_id
|
||||
|
||||
@property
|
||||
def prepaid_invoice(self):
|
||||
"""Get the prepaid invoice for this period"""
|
||||
return Invoice.query.filter_by(
|
||||
license_period_id=self.id,
|
||||
invoice_type=PaymentType.PREPAID
|
||||
).first()
|
||||
|
||||
@property
|
||||
def overage_invoice(self):
|
||||
"""Get the overage invoice for this period"""
|
||||
return Invoice.query.filter_by(
|
||||
license_period_id=self.id,
|
||||
invoice_type=PaymentType.POSTPAID
|
||||
).first()
|
||||
|
||||
@property
|
||||
def prepaid_payment(self):
|
||||
"""Get the prepaid payment for this period"""
|
||||
return Payment.query.filter_by(
|
||||
license_period_id=self.id,
|
||||
payment_type=PaymentType.PREPAID
|
||||
).first()
|
||||
|
||||
@property
|
||||
def overage_payment(self):
|
||||
"""Get the overage payment for this period"""
|
||||
return Payment.query.filter_by(
|
||||
license_period_id=self.id,
|
||||
payment_type=PaymentType.POSTPAID
|
||||
).first()
|
||||
|
||||
@property
|
||||
def all_invoices(self):
|
||||
"""Get all invoices for this period"""
|
||||
return self.invoices
|
||||
|
||||
@property
|
||||
def all_payments(self):
|
||||
"""Get all payments for this period"""
|
||||
return self.payments
|
||||
|
||||
def transition_status(self, new_status: PeriodStatus, user_id: int = None):
|
||||
"""Transition to a new status with proper validation and logging"""
|
||||
if not self.can_transition_to(new_status):
|
||||
raise ValueError(f"Invalid status transition from {self.status} to {new_status}")
|
||||
|
||||
self.status = new_status
|
||||
self.updated_at = dt.now(tz.utc)
|
||||
if user_id:
|
||||
self.updated_by = user_id
|
||||
|
||||
# Set appropriate timestamps
|
||||
if new_status == PeriodStatus.ACTIVE and not self.prepaid_received_at:
|
||||
self.prepaid_received_at = dt.now(tz.utc)
|
||||
elif new_status == PeriodStatus.COMPLETED:
|
||||
self.completed_at = dt.now(tz.utc)
|
||||
elif new_status == PeriodStatus.INVOICED:
|
||||
self.invoiced_at = dt.now(tz.utc)
|
||||
elif new_status == PeriodStatus.CLOSED:
|
||||
self.closed_at = dt.now(tz.utc)
|
||||
|
||||
@property
|
||||
def is_overdue(self):
|
||||
"""Check if a prepaid payment is overdue"""
|
||||
return (self.status == PeriodStatus.PENDING and
|
||||
self.period_start <= dt.now(tz.utc).date())
|
||||
|
||||
def can_transition_to(self, new_status: PeriodStatus) -> bool:
|
||||
"""Check if a status transition is valid"""
|
||||
valid_transitions = {
|
||||
PeriodStatus.UPCOMING: [PeriodStatus.ACTIVE, PeriodStatus.PENDING],
|
||||
PeriodStatus.PENDING: [PeriodStatus.ACTIVE],
|
||||
PeriodStatus.ACTIVE: [PeriodStatus.COMPLETED],
|
||||
PeriodStatus.COMPLETED: [PeriodStatus.INVOICED, PeriodStatus.CLOSED],
|
||||
PeriodStatus.INVOICED: [PeriodStatus.CLOSED],
|
||||
PeriodStatus.CLOSED: []
|
||||
}
|
||||
return new_status in valid_transitions.get(self.status, [])
|
||||
|
||||
def __repr__(self):
|
||||
return f'<LicensePeriod {self.id}: License {self.license_id}, Period {self.period_number}>'
|
||||
|
||||
|
||||
class LicenseUsage(db.Model):
|
||||
@@ -92,7 +350,6 @@ class LicenseUsage(db.Model):
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
license_id = db.Column(db.Integer, db.ForeignKey('public.license.id'), nullable=False)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
storage_mb_used = db.Column(db.Float, default=0)
|
||||
embedding_mb_used = db.Column(db.Float, default=0)
|
||||
@@ -102,9 +359,170 @@ class LicenseUsage(db.Model):
|
||||
interaction_prompt_tokens_used = db.Column(db.Integer, default=0)
|
||||
interaction_completion_tokens_used = db.Column(db.Integer, default=0)
|
||||
interaction_total_tokens_used = db.Column(db.Integer, default=0)
|
||||
period_start_date = db.Column(db.Date, nullable=False)
|
||||
period_end_date = db.Column(db.Date, nullable=False)
|
||||
license_period_id = db.Column(db.Integer, db.ForeignKey('public.license_period.id'), nullable=False)
|
||||
|
||||
# Standard audit fields
|
||||
created_at = db.Column(db.DateTime, server_default=db.func.now())
|
||||
updated_at = db.Column(db.DateTime, server_default=db.func.now(), onupdate=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
|
||||
license_period = db.relationship('LicensePeriod', back_populates='license_usage')
|
||||
|
||||
def recalculate_storage(self):
|
||||
Database(self.tenant_id).switch_schema()
|
||||
# Perform a SUM operation to get the total file size from document_versions
|
||||
total_storage = db.session.execute(text(f"""
|
||||
SELECT SUM(file_size)
|
||||
FROM document_version
|
||||
""")).scalar()
|
||||
|
||||
self.storage_mb_used = total_storage
|
||||
|
||||
|
||||
class PaymentType(Enum):
|
||||
PREPAID = "PREPAID"
|
||||
POSTPAID = "POSTPAID"
|
||||
|
||||
|
||||
class PaymentStatus(Enum):
|
||||
PENDING = "PENDING"
|
||||
PAID = "PAID"
|
||||
FAILED = "FAILED"
|
||||
CANCELLED = "CANCELLED"
|
||||
|
||||
|
||||
class Payment(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
license_period_id = db.Column(db.Integer, db.ForeignKey('public.license_period.id'), nullable=True)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
|
||||
# Payment details
|
||||
payment_type = db.Column(db.Enum(PaymentType), nullable=False)
|
||||
amount = db.Column(db.Numeric(10, 2), nullable=False)
|
||||
currency = db.Column(db.String(3), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
|
||||
# Status tracking
|
||||
status = db.Column(db.Enum(PaymentStatus), nullable=False, default=PaymentStatus.PENDING)
|
||||
|
||||
# External provider information
|
||||
external_payment_id = db.Column(db.String(255), nullable=True)
|
||||
payment_method = db.Column(db.String(50), nullable=True) # credit_card, bank_transfer, etc.
|
||||
provider_data = db.Column(JSONB, nullable=True) # Provider-specific data
|
||||
|
||||
# Payment information
|
||||
paid_at = db.Column(db.DateTime, nullable=True)
|
||||
|
||||
# Standard audit fields
|
||||
created_at = db.Column(db.DateTime, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
updated_at = db.Column(db.DateTime, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
|
||||
# Relationships
|
||||
license_period = db.relationship('LicensePeriod', back_populates='payments')
|
||||
invoice = db.relationship('Invoice', back_populates='payment', uselist=False)
|
||||
|
||||
@property
|
||||
def is_overdue(self):
|
||||
"""Check if payment is overdue"""
|
||||
if self.status != PaymentStatus.PENDING:
|
||||
return False
|
||||
|
||||
# For prepaid payments, check if period start has passed
|
||||
if (self.payment_type == PaymentType.PREPAID and
|
||||
self.license_period_id):
|
||||
return self.license_period.period_start <= dt.now(tz.utc).date()
|
||||
|
||||
# For postpaid, check against due date (would be on invoice)
|
||||
return False
|
||||
|
||||
def __repr__(self):
|
||||
return f'<Payment {self.id}: {self.payment_type} {self.amount} {self.currency}>'
|
||||
|
||||
|
||||
class InvoiceStatus(Enum):
|
||||
DRAFT = "DRAFT"
|
||||
SENT = "SENT"
|
||||
PAID = "PAID"
|
||||
OVERDUE = "OVERDUE"
|
||||
CANCELLED = "CANCELLED"
|
||||
|
||||
|
||||
class Invoice(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
license_period_id = db.Column(db.Integer, db.ForeignKey('public.license_period.id'), nullable=False)
|
||||
payment_id = db.Column(db.Integer, db.ForeignKey('public.payment.id'), nullable=True)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
|
||||
# Invoice details
|
||||
invoice_type = db.Column(db.Enum(PaymentType), nullable=False)
|
||||
invoice_number = db.Column(db.String(50), unique=True, nullable=False)
|
||||
invoice_date = db.Column(db.Date, nullable=False)
|
||||
due_date = db.Column(db.Date, nullable=False)
|
||||
|
||||
# Financial details
|
||||
amount = db.Column(db.Numeric(10, 2), nullable=False)
|
||||
currency = db.Column(db.String(3), nullable=False)
|
||||
tax_amount = db.Column(db.Numeric(10, 2), default=0)
|
||||
|
||||
# Descriptive fields
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
status = db.Column(db.Enum(InvoiceStatus), nullable=False, default=InvoiceStatus.DRAFT)
|
||||
|
||||
# Timestamps
|
||||
sent_at = db.Column(db.DateTime, nullable=True)
|
||||
paid_at = db.Column(db.DateTime, nullable=True)
|
||||
|
||||
# Standard audit fields
|
||||
created_at = db.Column(db.DateTime, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
updated_at = db.Column(db.DateTime, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
|
||||
# Relationships
|
||||
license_period = db.relationship('LicensePeriod', back_populates='invoices')
|
||||
payment = db.relationship('Payment', back_populates='invoice')
|
||||
|
||||
def __repr__(self):
|
||||
return f'<Invoice {self.invoice_number}: {self.amount} {self.currency}>'
|
||||
|
||||
|
||||
class LicenseChangeLog(db.Model):
|
||||
"""
|
||||
Log of changes to license configurations
|
||||
Used for auditing and tracking when/why license details changed
|
||||
"""
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
license_id = db.Column(db.Integer, db.ForeignKey('public.license.id'), nullable=False)
|
||||
changed_at = db.Column(db.DateTime, nullable=False, default=lambda: dt.now(tz.utc))
|
||||
|
||||
# What changed
|
||||
field_name = db.Column(db.String(100), nullable=False)
|
||||
old_value = db.Column(db.String(255), nullable=True)
|
||||
new_value = db.Column(db.String(255), nullable=False)
|
||||
|
||||
# Why it changed
|
||||
reason = db.Column(db.Text, nullable=True)
|
||||
|
||||
# Standard audit fields
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
# Relationships
|
||||
license = db.relationship('License', backref=db.backref('change_logs', order_by='LicenseChangeLog.changed_at'))
|
||||
|
||||
def __repr__(self):
|
||||
return f'<LicenseChangeLog: {self.license_id} {self.field_name} {self.old_value} -> {self.new_value}>'
|
||||
|
||||
license = db.relationship('License', back_populates='usages')
|
||||
|
||||
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
from ..extensions import db
|
||||
from .user import User, Tenant
|
||||
from .document import Embedding
|
||||
from .document import Embedding, Retriever
|
||||
|
||||
|
||||
class ChatSession(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
user_id = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
session_id = db.Column(db.String(36), nullable=True)
|
||||
session_id = db.Column(db.String(49), nullable=True)
|
||||
session_start = db.Column(db.DateTime, nullable=False)
|
||||
session_end = db.Column(db.DateTime, nullable=True)
|
||||
timezone = db.Column(db.String(30), nullable=True)
|
||||
@@ -18,14 +20,168 @@ class ChatSession(db.Model):
|
||||
return f"<ChatSession {self.id} by {self.user_id}>"
|
||||
|
||||
|
||||
class Specialist(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_RAG")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="1.0.0")
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Relationship to retrievers through the association table
|
||||
retrievers = db.relationship('SpecialistRetriever', backref='specialist', lazy=True,
|
||||
cascade="all, delete-orphan")
|
||||
agents = db.relationship('EveAIAgent', backref='specialist', lazy=True)
|
||||
tasks = db.relationship('EveAITask', backref='specialist', lazy=True)
|
||||
tools = db.relationship('EveAITool', backref='specialist', lazy=True)
|
||||
dispatchers = db.relationship('SpecialistDispatcher', backref='specialist', lazy=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
|
||||
class EveAIAsset(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="DOCUMENT_TEMPLATE")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="1.0.0")
|
||||
valid_from = db.Column(db.DateTime, nullable=True)
|
||||
valid_to = db.Column(db.DateTime, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
# Relations
|
||||
versions = db.relationship('EveAIAssetVersion', backref='asset', lazy=True)
|
||||
|
||||
|
||||
class EveAIAssetVersion(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
asset_id = db.Column(db.Integer, db.ForeignKey(EveAIAsset.id), nullable=False)
|
||||
bucket_name = db.Column(db.String(255), nullable=True)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
# Relations
|
||||
instructions = db.relationship('EveAIAssetInstruction', backref='asset_version', lazy=True)
|
||||
|
||||
|
||||
class EveAIAssetInstruction(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
asset_version_id = db.Column(db.Integer, db.ForeignKey(EveAIAssetVersion.id), nullable=False)
|
||||
name = db.Column(db.String(255), nullable=False)
|
||||
content = db.Column(db.Text, nullable=True)
|
||||
|
||||
|
||||
class EveAIProcessedAsset(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
asset_version_id = db.Column(db.Integer, db.ForeignKey(EveAIAssetVersion.id), nullable=False)
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id), nullable=True)
|
||||
chat_session_id = db.Column(db.Integer, db.ForeignKey(ChatSession.id), nullable=True)
|
||||
bucket_name = db.Column(db.String(255), nullable=True)
|
||||
object_name = db.Column(db.String(255), nullable=True)
|
||||
created_at = db.Column(db.DateTime, nullable=True, server_default=db.func.now())
|
||||
|
||||
|
||||
class EveAIAgent(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id), nullable=False)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_RAG")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="1.0.0")
|
||||
role = db.Column(db.Text, nullable=True)
|
||||
goal = db.Column(db.Text, nullable=True)
|
||||
backstory = db.Column(db.Text, nullable=True)
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
|
||||
class EveAITask(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id), nullable=False)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_RAG")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="1.0.0")
|
||||
task_description = db.Column(db.Text, nullable=True)
|
||||
expected_output = db.Column(db.Text, nullable=True)
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
context = db.Column(JSONB, nullable=True)
|
||||
asynchronous = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
|
||||
class EveAITool(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id), nullable=False)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_RAG")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="1.0.0")
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
|
||||
class Dispatcher(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(50), nullable=False, default="STANDARD_RAG")
|
||||
type_version = db.Column(db.String(20), nullable=True, default="1.0.0")
|
||||
tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
configuration = db.Column(JSONB, nullable=True)
|
||||
arguments = db.Column(JSONB, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
|
||||
|
||||
class Interaction(db.Model):
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
chat_session_id = db.Column(db.Integer, db.ForeignKey(ChatSession.id), nullable=False)
|
||||
question = db.Column(db.Text, nullable=False)
|
||||
detailed_question = db.Column(db.Text, nullable=True)
|
||||
answer = db.Column(db.Text, nullable=True)
|
||||
algorithm_used = db.Column(db.String(20), nullable=True)
|
||||
language = db.Column(db.String(2), nullable=False)
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id), nullable=True)
|
||||
specialist_arguments = db.Column(JSONB, nullable=True)
|
||||
specialist_results = db.Column(JSONB, nullable=True)
|
||||
timezone = db.Column(db.String(30), nullable=True)
|
||||
appreciation = db.Column(db.Integer, nullable=True)
|
||||
|
||||
@@ -33,6 +189,7 @@ class Interaction(db.Model):
|
||||
question_at = db.Column(db.DateTime, nullable=False)
|
||||
detailed_question_at = db.Column(db.DateTime, nullable=True)
|
||||
answer_at = db.Column(db.DateTime, nullable=True)
|
||||
processing_error = db.Column(db.String(255), nullable=True)
|
||||
|
||||
# Relations
|
||||
embeddings = db.relationship('InteractionEmbedding', backref='interaction', lazy=True)
|
||||
@@ -44,3 +201,17 @@ class Interaction(db.Model):
|
||||
class InteractionEmbedding(db.Model):
|
||||
interaction_id = db.Column(db.Integer, db.ForeignKey(Interaction.id, ondelete='CASCADE'), primary_key=True)
|
||||
embedding_id = db.Column(db.Integer, db.ForeignKey(Embedding.id, ondelete='CASCADE'), primary_key=True)
|
||||
|
||||
|
||||
class SpecialistRetriever(db.Model):
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id, ondelete='CASCADE'), primary_key=True)
|
||||
retriever_id = db.Column(db.Integer, db.ForeignKey(Retriever.id, ondelete='CASCADE'), primary_key=True)
|
||||
|
||||
retriever = db.relationship("Retriever", backref="specialist_retrievers")
|
||||
|
||||
|
||||
class SpecialistDispatcher(db.Model):
|
||||
specialist_id = db.Column(db.Integer, db.ForeignKey(Specialist.id, ondelete='CASCADE'), primary_key=True)
|
||||
dispatcher_id = db.Column(db.Integer, db.ForeignKey(Dispatcher.id, ondelete='CASCADE'), primary_key=True)
|
||||
|
||||
dispatcher = db.relationship("Dispatcher", backref="specialist_dispatchers")
|
||||
|
||||
@@ -20,50 +20,18 @@ class Tenant(db.Model):
|
||||
|
||||
# company Information
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
code = db.Column(db.String(50), unique=True, nullable=True)
|
||||
name = db.Column(db.String(80), unique=True, nullable=False)
|
||||
website = db.Column(db.String(255), nullable=True)
|
||||
timezone = db.Column(db.String(50), nullable=True, default='UTC')
|
||||
rag_context = db.Column(db.Text, nullable=True)
|
||||
type = db.Column(db.String(20), nullable=True, server_default='Active')
|
||||
|
||||
# language information
|
||||
default_language = db.Column(db.String(2), nullable=True)
|
||||
allowed_languages = db.Column(ARRAY(sa.String(2)), nullable=True)
|
||||
|
||||
# LLM specific choices
|
||||
embedding_model = db.Column(db.String(50), nullable=True)
|
||||
llm_model = db.Column(db.String(50), nullable=True)
|
||||
|
||||
# # Embedding variables ==> To be removed once all migrations (dev + prod) have been done
|
||||
# html_tags = db.Column(ARRAY(sa.String(10)), nullable=True, default=['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'li'])
|
||||
# html_end_tags = db.Column(ARRAY(sa.String(10)), nullable=True, default=['p', 'li'])
|
||||
# html_included_elements = db.Column(ARRAY(sa.String(50)), nullable=True)
|
||||
# html_excluded_elements = db.Column(ARRAY(sa.String(50)), nullable=True)
|
||||
# html_excluded_classes = db.Column(ARRAY(sa.String(200)), nullable=True)
|
||||
#
|
||||
# min_chunk_size = db.Column(db.Integer, nullable=True, default=2000)
|
||||
# max_chunk_size = db.Column(db.Integer, nullable=True, default=3000)
|
||||
#
|
||||
# # Embedding search variables
|
||||
# es_k = db.Column(db.Integer, nullable=True, default=5)
|
||||
# es_similarity_threshold = db.Column(db.Float, nullable=True, default=0.7)
|
||||
#
|
||||
# # Chat variables
|
||||
# chat_RAG_temperature = db.Column(db.Float, nullable=True, default=0.3)
|
||||
# chat_no_RAG_temperature = db.Column(db.Float, nullable=True, default=0.5)
|
||||
fallback_algorithms = db.Column(ARRAY(sa.String(50)), nullable=True)
|
||||
|
||||
# Licensing Information
|
||||
encrypted_chat_api_key = db.Column(db.String(500), nullable=True)
|
||||
encrypted_api_key = db.Column(db.String(500), nullable=True)
|
||||
|
||||
# # Tuning enablers
|
||||
# embed_tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
# rag_tuning = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Entitlements
|
||||
# Entitlements
|
||||
currency = db.Column(db.String(20), nullable=True)
|
||||
usage_email = db.Column(db.String(255), nullable=True)
|
||||
storage_dirty = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Relations
|
||||
@@ -90,15 +58,10 @@ class Tenant(db.Model):
|
||||
'name': self.name,
|
||||
'website': self.website,
|
||||
'timezone': self.timezone,
|
||||
'rag_context': self.rag_context,
|
||||
'type': self.type,
|
||||
'default_language': self.default_language,
|
||||
'allowed_languages': self.allowed_languages,
|
||||
'embedding_model': self.embedding_model,
|
||||
'llm_model': self.llm_model,
|
||||
'fallback_algorithms': self.fallback_algorithms,
|
||||
'currency': self.currency,
|
||||
'usage_email': self.usage_email,
|
||||
}
|
||||
|
||||
|
||||
@@ -131,6 +94,7 @@ class User(db.Model, UserMixin):
|
||||
|
||||
# User Information
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
user_name = db.Column(db.String(80), unique=True, nullable=False)
|
||||
email = db.Column(db.String(255), unique=True, nullable=False)
|
||||
password = db.Column(db.String(255), nullable=True)
|
||||
@@ -140,6 +104,8 @@ class User(db.Model, UserMixin):
|
||||
fs_uniquifier = db.Column(db.String(255), unique=True, nullable=False)
|
||||
confirmed_at = db.Column(db.DateTime, nullable=True)
|
||||
valid_to = db.Column(db.Date, nullable=True)
|
||||
is_primary_contact = db.Column(db.Boolean, nullable=True, default=False)
|
||||
is_financial_contact = db.Column(db.Boolean, nullable=True, default=False)
|
||||
|
||||
# Security Trackable Information
|
||||
last_login_at = db.Column(db.DateTime, nullable=True)
|
||||
@@ -150,7 +116,6 @@ class User(db.Model, UserMixin):
|
||||
|
||||
# Relations
|
||||
roles = db.relationship('Role', secondary=RolesUsers.__table__, backref=db.backref('users', lazy='dynamic'))
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
|
||||
def __repr__(self):
|
||||
return '<User %r>' % self.user_name
|
||||
@@ -173,10 +138,136 @@ class TenantDomain(db.Model):
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey(User.id), nullable=False)
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=False)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey(User.id))
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
|
||||
def __repr__(self):
|
||||
return f"<TenantDomain {self.id}: {self.domain}>"
|
||||
|
||||
|
||||
class TenantProject(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False)
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
services = db.Column(ARRAY(sa.String(50)), nullable=False)
|
||||
encrypted_api_key = db.Column(db.String(500), nullable=True)
|
||||
visual_api_key = db.Column(db.String(20), nullable=True)
|
||||
active = db.Column(db.Boolean, nullable=False, default=True)
|
||||
responsible_email = db.Column(db.String(255), nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'))
|
||||
|
||||
# Relations
|
||||
tenant = db.relationship('Tenant', backref='projects')
|
||||
|
||||
def __repr__(self):
|
||||
return f"<TenantProject {self.id}: {self.name}>"
|
||||
|
||||
|
||||
class Partner(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), nullable=False, unique=True)
|
||||
code = db.Column(db.String(50), unique=True, nullable=False)
|
||||
|
||||
# Basic information
|
||||
logo_url = db.Column(db.String(255), nullable=True)
|
||||
active = db.Column(db.Boolean, default=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
# Relationships
|
||||
services = db.relationship('PartnerService', back_populates='partner')
|
||||
tenant = db.relationship('Tenant', backref=db.backref('partner', uselist=False))
|
||||
|
||||
def to_dict(self):
|
||||
services_info = []
|
||||
for service in self.services:
|
||||
services_info.append({
|
||||
'id': service.id,
|
||||
'name': service.name,
|
||||
'description': service.description,
|
||||
'type': service.type,
|
||||
'type_version': service.type_version,
|
||||
'active': service.active,
|
||||
'configuration': service.configuration,
|
||||
'permissions': service.permissions,
|
||||
})
|
||||
return {
|
||||
'id': self.id,
|
||||
'tenant_id': self.tenant_id,
|
||||
'code': self.code,
|
||||
'logo_url': self.logo_url,
|
||||
'active': self.active,
|
||||
'name': self.tenant.name,
|
||||
'services': services_info
|
||||
}
|
||||
|
||||
|
||||
class PartnerService(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
id = db.Column(db.Integer, primary_key=True)
|
||||
partner_id = db.Column(db.Integer, db.ForeignKey('public.partner.id'), nullable=False)
|
||||
|
||||
# Basic info
|
||||
name = db.Column(db.String(50), nullable=False)
|
||||
description = db.Column(db.Text, nullable=True)
|
||||
|
||||
# Service type with versioning (similar to your specialist/retriever pattern)
|
||||
type = db.Column(db.String(50), nullable=False) # REFERRAL, KNOWLEDGE, SPECIALIST, IMPLEMENTATION, WHITE_LABEL
|
||||
type_version = db.Column(db.String(20), nullable=False, default="1.0.0")
|
||||
|
||||
# Status
|
||||
active = db.Column(db.Boolean, default=True)
|
||||
|
||||
# Dynamic configuration specific to this service - using JSONB like your other models
|
||||
configuration = db.Column(db.JSON, nullable=True)
|
||||
permissions = db.Column(db.JSON, nullable=True)
|
||||
|
||||
# For services that need to track shared resources
|
||||
system_metadata = db.Column(db.JSON, nullable=True)
|
||||
user_metadata = db.Column(db.JSON, nullable=True)
|
||||
|
||||
# Versioning Information
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
# Relationships
|
||||
partner = db.relationship('Partner', back_populates='services')
|
||||
license_tiers = db.relationship('PartnerServiceLicenseTier', back_populates='partner_service')
|
||||
|
||||
|
||||
class PartnerTenant(db.Model):
|
||||
__bind_key__ = 'public'
|
||||
__table_args__ = {'schema': 'public'}
|
||||
|
||||
partner_service_id = db.Column(db.Integer, db.ForeignKey('public.partner_service.id'), primary_key=True)
|
||||
tenant_id = db.Column(db.Integer, db.ForeignKey('public.tenant.id'), primary_key=True)
|
||||
|
||||
# JSONB for flexible configuration specific to this relationship
|
||||
configuration = db.Column(db.JSON, nullable=True)
|
||||
|
||||
# Tracking
|
||||
created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now())
|
||||
created_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
updated_at = db.Column(db.DateTime, nullable=False, server_default=db.func.now(), onupdate=db.func.now())
|
||||
updated_by = db.Column(db.Integer, db.ForeignKey('public.user.id'), nullable=True)
|
||||
|
||||
9
common/services/entitlements/__init__.py
Normal file
9
common/services/entitlements/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
from common.services.entitlements.license_period_services import LicensePeriodServices
|
||||
from common.services.entitlements.license_usage_services import LicenseUsageServices
|
||||
from common.services.entitlements.license_tier_services import LicenseTierServices
|
||||
|
||||
__all__ = [
|
||||
'LicensePeriodServices',
|
||||
'LicenseUsageServices',
|
||||
'LicenseTierServices'
|
||||
]
|
||||
247
common/services/entitlements/license_period_services.py
Normal file
247
common/services/entitlements/license_period_services.py
Normal file
@@ -0,0 +1,247 @@
|
||||
from dateutil.relativedelta import relativedelta
|
||||
from datetime import datetime as dt, timezone as tz, timedelta
|
||||
from flask import current_app
|
||||
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from sqlalchemy.sql.expression import and_
|
||||
|
||||
from common.extensions import db
|
||||
from common.models.entitlements import LicensePeriod, License, PeriodStatus, LicenseUsage
|
||||
from common.utils.eveai_exceptions import EveAILicensePeriodsExceeded, EveAIPendingLicensePeriod, EveAINoActiveLicense
|
||||
from common.utils.model_logging_utils import set_logging_information, update_logging_information
|
||||
|
||||
|
||||
class LicensePeriodServices:
|
||||
@staticmethod
|
||||
def find_current_license_period_for_usage(tenant_id: int) -> LicensePeriod:
|
||||
"""
|
||||
Find the current license period for a tenant. It ensures the status of the different license periods are adapted
|
||||
when required, and a LicenseUsage object is created if required.
|
||||
|
||||
Args:
|
||||
tenant_id: The ID of the tenant to find the license period for
|
||||
|
||||
Raises:
|
||||
EveAIException: and derived classes
|
||||
"""
|
||||
try:
|
||||
current_app.logger.debug(f"Finding current license period for tenant {tenant_id}")
|
||||
current_date = dt.now(tz.utc).date()
|
||||
license_period = (db.session.query(LicensePeriod)
|
||||
.filter_by(tenant_id=tenant_id)
|
||||
.filter(and_(LicensePeriod.period_start <= current_date,
|
||||
LicensePeriod.period_end >= current_date))
|
||||
.first())
|
||||
current_app.logger.debug(f"End searching for license period for tenant {tenant_id} ")
|
||||
if not license_period:
|
||||
current_app.logger.debug(f"No license period found for tenant {tenant_id} on date {current_date}")
|
||||
license_period = LicensePeriodServices._create_next_license_period_for_usage(tenant_id)
|
||||
current_app.logger.debug(f"Created license period {license_period.id} for tenant {tenant_id}")
|
||||
if license_period:
|
||||
current_app.logger.debug(f"Found license period {license_period.id} for tenant {tenant_id} "
|
||||
f"with status {license_period.status}")
|
||||
match license_period.status:
|
||||
case PeriodStatus.UPCOMING:
|
||||
current_app.logger.debug(f"In upcoming state")
|
||||
LicensePeriodServices._complete_last_license_period(tenant_id=tenant_id)
|
||||
current_app.logger.debug(f"Completed last license period for tenant {tenant_id}")
|
||||
LicensePeriodServices._activate_license_period(license_period=license_period)
|
||||
current_app.logger.debug(f"Activated license period {license_period.id} for tenant {tenant_id}")
|
||||
if not license_period.license_usage:
|
||||
new_license_usage = LicenseUsage(
|
||||
tenant_id=tenant_id,
|
||||
)
|
||||
new_license_usage.license_period = license_period
|
||||
try:
|
||||
db.session.add(new_license_usage)
|
||||
db.session.commit()
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(
|
||||
f"Error creating new license usage for license period "
|
||||
f"{license_period.id}: {str(e)}")
|
||||
raise e
|
||||
if license_period.status == PeriodStatus.ACTIVE:
|
||||
return license_period
|
||||
else:
|
||||
# Status is PENDING, so no prepaid payment received. There is no license period we can use.
|
||||
# We allow for a delay of 5 days before raising an exception.
|
||||
current_date = dt.now(tz.utc).date()
|
||||
delta = abs(current_date - license_period.period_start)
|
||||
if delta > timedelta(days=current_app.config.get('ENTITLEMENTS_MAX_PENDING_DAYS', 5)):
|
||||
raise EveAIPendingLicensePeriod()
|
||||
case PeriodStatus.ACTIVE:
|
||||
return license_period
|
||||
case PeriodStatus.PENDING:
|
||||
return license_period
|
||||
else:
|
||||
raise EveAILicensePeriodsExceeded(license_id=None)
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error finding current license period for tenant {tenant_id}: {str(e)}")
|
||||
raise e
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def _create_next_license_period_for_usage(tenant_id) -> LicensePeriod:
|
||||
"""
|
||||
Create a new period for this license using the current license configuration
|
||||
|
||||
Args:
|
||||
tenant_id: The ID of the tenant to create the period for
|
||||
|
||||
Returns:
|
||||
LicensePeriod: The newly created license period
|
||||
"""
|
||||
current_date = dt.now(tz.utc).date()
|
||||
|
||||
# Zoek de actieve licentie voor deze tenant op de huidige datum
|
||||
the_license = (db.session.query(License)
|
||||
.filter_by(tenant_id=tenant_id)
|
||||
.filter(License.start_date <= current_date)
|
||||
.filter(License.end_date >= current_date)
|
||||
.first())
|
||||
|
||||
if not the_license:
|
||||
current_app.logger.error(f"No active license found for tenant {tenant_id} on date {current_date}")
|
||||
raise EveAINoActiveLicense(tenant_id=tenant_id)
|
||||
else:
|
||||
current_app.logger.debug(f"Found active license {the_license.id} for tenant {tenant_id} "
|
||||
f"on date {current_date}")
|
||||
|
||||
next_period_number = 1
|
||||
if the_license.periods:
|
||||
# If there are existing periods, get the next sequential number
|
||||
next_period_number = max(p.period_number for p in the_license.periods) + 1
|
||||
current_app.logger.debug(f"Next period number for tenant {tenant_id} is {next_period_number}")
|
||||
|
||||
if next_period_number > the_license.nr_of_periods:
|
||||
raise EveAILicensePeriodsExceeded(license_id=the_license.id)
|
||||
|
||||
new_license_period = LicensePeriod(
|
||||
license_id=the_license.id,
|
||||
tenant_id=tenant_id,
|
||||
period_number=next_period_number,
|
||||
period_start=the_license.start_date + relativedelta(months=next_period_number-1),
|
||||
period_end=the_license.end_date + relativedelta(months=next_period_number, days=-1),
|
||||
status=PeriodStatus.UPCOMING,
|
||||
upcoming_at=dt.now(tz.utc),
|
||||
)
|
||||
set_logging_information(new_license_period, dt.now(tz.utc))
|
||||
|
||||
try:
|
||||
current_app.logger.debug(f"Creating next license period for tenant {tenant_id} ")
|
||||
db.session.add(new_license_period)
|
||||
db.session.commit()
|
||||
current_app.logger.info(f"Created next license period for tenant {tenant_id} "
|
||||
f"with id {new_license_period.id}")
|
||||
return new_license_period
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error creating next license period for tenant {tenant_id}: {str(e)}")
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def _activate_license_period(license_period_id: int = None, license_period: LicensePeriod = None) -> LicensePeriod:
|
||||
"""
|
||||
Activate a license period
|
||||
|
||||
Args:
|
||||
license_period_id: The ID of the license period to activate (optional if license_period is provided)
|
||||
license_period: The LicensePeriod object to activate (optional if license_period_id is provided)
|
||||
|
||||
Returns:
|
||||
LicensePeriod: The activated license period object
|
||||
|
||||
Raises:
|
||||
ValueError: If neither license_period_id nor license_period is provided
|
||||
"""
|
||||
current_app.logger.debug(f"Activating license period")
|
||||
if license_period is None and license_period_id is None:
|
||||
raise ValueError("Either license_period_id or license_period must be provided")
|
||||
|
||||
# Get a license period object if only ID was provided
|
||||
if license_period is None:
|
||||
current_app.logger.debug(f"Getting license period {license_period_id} to activate")
|
||||
license_period = LicensePeriod.query.get_or_404(license_period_id)
|
||||
|
||||
if license_period.pending_at is not None:
|
||||
license_period.pending_at = dt.now(tz.utc)
|
||||
license_period.status = PeriodStatus.PENDING
|
||||
if license_period.prepaid_payment:
|
||||
# There is a payment received for the given period
|
||||
license_period.active_at = dt.now(tz.utc)
|
||||
license_period.status = PeriodStatus.ACTIVE
|
||||
|
||||
# Copy snapshot fields from the license to the period
|
||||
the_license = License.query.get_or_404(license_period.license_id)
|
||||
license_period.currency = the_license.currency
|
||||
license_period.basic_fee = the_license.basic_fee
|
||||
license_period.max_storage_mb = the_license.max_storage_mb
|
||||
license_period.additional_storage_price = the_license.additional_storage_price
|
||||
license_period.additional_storage_bucket = the_license.additional_storage_bucket
|
||||
license_period.included_embedding_mb = the_license.included_embedding_mb
|
||||
license_period.additional_embedding_price = the_license.additional_embedding_price
|
||||
license_period.additional_embedding_bucket = the_license.additional_embedding_bucket
|
||||
license_period.included_interaction_tokens = the_license.included_interaction_tokens
|
||||
license_period.additional_interaction_token_price = the_license.additional_interaction_token_price
|
||||
license_period.additional_interaction_bucket = the_license.additional_interaction_bucket
|
||||
license_period.additional_storage_allowed = the_license.additional_storage_allowed
|
||||
license_period.additional_embedding_allowed = the_license.additional_embedding_allowed
|
||||
license_period.additional_interaction_allowed = the_license.additional_interaction_allowed
|
||||
|
||||
update_logging_information(license_period, dt.now(tz.utc))
|
||||
|
||||
if not license_period.license_usage:
|
||||
license_period.license_usage = LicenseUsage(
|
||||
tenant_id=license_period.tenant_id,
|
||||
license_period_id=license_period.id,
|
||||
)
|
||||
|
||||
license_period.license_usage.recalculate_storage()
|
||||
|
||||
try:
|
||||
db.session.add(license_period)
|
||||
db.session.add(license_period.license_usage)
|
||||
db.session.commit()
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error activating license period {license_period_id}: {str(e)}")
|
||||
raise e
|
||||
|
||||
return license_period
|
||||
|
||||
@staticmethod
|
||||
def _complete_last_license_period(tenant_id) -> None:
|
||||
"""
|
||||
Complete the active or pending license period for a tenant. This is done by setting the status to COMPLETED.
|
||||
|
||||
Args:
|
||||
tenant_id: De ID van de tenant
|
||||
"""
|
||||
# Zoek de licenseperiode voor deze tenant met status ACTIVE of PENDING
|
||||
active_period = (
|
||||
db.session.query(LicensePeriod)
|
||||
.filter_by(tenant_id=tenant_id)
|
||||
.filter(LicensePeriod.status.in_([PeriodStatus.ACTIVE, PeriodStatus.PENDING]))
|
||||
.first()
|
||||
)
|
||||
|
||||
# Als er geen actieve periode gevonden is, hoeven we niets te doen
|
||||
if not active_period:
|
||||
return
|
||||
|
||||
# Zet de gevonden periode op COMPLETED
|
||||
active_period.status = PeriodStatus.COMPLETED
|
||||
active_period.completed_at = dt.now(tz.utc)
|
||||
update_logging_information(active_period, dt.now(tz.utc))
|
||||
|
||||
try:
|
||||
db.session.add(active_period)
|
||||
db.session.commit()
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error completing period {active_period.id} for {tenant_id}: {str(e)}")
|
||||
raise e
|
||||
67
common/services/entitlements/license_tier_services.py
Normal file
67
common/services/entitlements/license_tier_services.py
Normal file
@@ -0,0 +1,67 @@
|
||||
from flask import session, flash, current_app
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from common.extensions import db
|
||||
from common.models.entitlements import PartnerServiceLicenseTier
|
||||
from common.models.user import Partner
|
||||
from common.utils.eveai_exceptions import EveAINoManagementPartnerService
|
||||
from common.utils.model_logging_utils import set_logging_information
|
||||
|
||||
|
||||
class LicenseTierServices:
|
||||
@staticmethod
|
||||
def associate_license_tier_with_partner(license_tier_id):
|
||||
"""Associate a license tier with a partner"""
|
||||
try:
|
||||
partner_id = session['partner']['id']
|
||||
# Get partner service (MANAGEMENT_SERVICE type)
|
||||
partner = Partner.query.get(partner_id)
|
||||
if not partner:
|
||||
return
|
||||
|
||||
# Find a management service for this partner
|
||||
management_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
|
||||
if not management_service:
|
||||
flash("Cannot associate license tier with partner. No management service defined for partner", "danger")
|
||||
current_app.logger.error(f"No Management Service defined for partner {partner_id}"
|
||||
f"trying to associate license tier {license_tier_id}.")
|
||||
raise EveAINoManagementPartnerService()
|
||||
# Check if the association already exists
|
||||
existing_association = PartnerServiceLicenseTier.query.filter_by(
|
||||
partner_service_id=management_service['id'],
|
||||
license_tier_id=license_tier_id
|
||||
).first()
|
||||
|
||||
if existing_association:
|
||||
# Association already exists, nothing to do
|
||||
flash("License tier was already associated with partner", "info")
|
||||
current_app.logger.info(f"Association between partner service {management_service['id']} and "
|
||||
f"license tier {license_tier_id} already exists.")
|
||||
return
|
||||
|
||||
# Create the association
|
||||
association = PartnerServiceLicenseTier(
|
||||
partner_service_id=management_service['id'],
|
||||
license_tier_id=license_tier_id
|
||||
)
|
||||
set_logging_information(association, dt.now(tz.utc))
|
||||
|
||||
db.session.add(association)
|
||||
db.session.commit()
|
||||
|
||||
flash("Successfully associated license tier to partner", "success")
|
||||
current_app.logger.info(f"Successfully associated license tier {license_tier_id} with "
|
||||
f"partner service {management_service['id']}")
|
||||
|
||||
return True
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
flash("Failed to associated license tier with partner service due to an internal error. "
|
||||
"Please contact the System Administrator", "danger")
|
||||
current_app.logger.error(f"Error associating license tier {license_tier_id} with partner: {str(e)}")
|
||||
raise e
|
||||
143
common/services/entitlements/license_usage_services.py
Normal file
143
common/services/entitlements/license_usage_services.py
Normal file
@@ -0,0 +1,143 @@
|
||||
from dateutil.relativedelta import relativedelta
|
||||
from flask import session, current_app, flash
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from sqlalchemy.sql.expression import text
|
||||
|
||||
from common.extensions import db, cache_manager
|
||||
from common.models.entitlements import PartnerServiceLicenseTier, License, LicenseUsage, LicensePeriod, PeriodStatus
|
||||
from common.models.user import Partner, PartnerTenant
|
||||
from common.services.entitlements import LicensePeriodServices
|
||||
from common.utils.database import Database
|
||||
from common.utils.eveai_exceptions import EveAINoManagementPartnerService, EveAINoActiveLicense, \
|
||||
EveAIStorageQuotaExceeded, EveAIEmbeddingQuotaExceeded, EveAIInteractionQuotaExceeded, EveAILicensePeriodsExceeded, \
|
||||
EveAIException
|
||||
from common.utils.model_logging_utils import set_logging_information, update_logging_information
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
|
||||
from common.utils.security_utils import current_user_has_role
|
||||
|
||||
|
||||
class LicenseUsageServices:
|
||||
@staticmethod
|
||||
def check_storage_and_embedding_quota(tenant_id: int, file_size_mb: float) -> None:
|
||||
"""
|
||||
Check if a tenant can add a new document without exceeding storage and embedding quotas
|
||||
|
||||
Args:
|
||||
tenant_id: ID of the tenant
|
||||
file_size_mb: Size of the file in MB
|
||||
|
||||
Raises:
|
||||
EveAIStorageQuotaExceeded: If storage quota would be exceeded
|
||||
EveAIEmbeddingQuotaExceeded: If embedding quota would be exceeded
|
||||
EveAINoActiveLicense: If no active license is found
|
||||
EveAIException: For other errors
|
||||
"""
|
||||
# Get active license period
|
||||
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
||||
# Early return if both overruns are allowed - no need to check usage at all
|
||||
if license_period.additional_storage_allowed and license_period.additional_embedding_allowed:
|
||||
return
|
||||
|
||||
# Check storage quota only if overruns are not allowed
|
||||
if not license_period.additional_storage_allowed:
|
||||
LicenseUsageServices._validate_storage_quota(license_period, file_size_mb)
|
||||
|
||||
# Check embedding quota only if overruns are not allowed
|
||||
if not license_period.additional_embedding_allowed:
|
||||
LicenseUsageServices._validate_embedding_quota(license_period, file_size_mb)
|
||||
|
||||
@staticmethod
|
||||
def check_embedding_quota(tenant_id: int, file_size_mb: float) -> None:
|
||||
"""
|
||||
Check if a tenant can re-embed a document without exceeding embedding quota
|
||||
|
||||
Args:
|
||||
tenant_id: ID of the tenant
|
||||
file_size_mb: Size of the file in MB
|
||||
|
||||
Raises:
|
||||
EveAIEmbeddingQuotaExceeded: If embedding quota would be exceeded
|
||||
EveAINoActiveLicense: If no active license is found
|
||||
EveAIException: For other errors
|
||||
"""
|
||||
# Get active license period
|
||||
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
||||
# Early return if both overruns are allowed - no need to check usage at all
|
||||
if license_period.additional_embedding_allowed:
|
||||
return
|
||||
|
||||
# Check embedding quota
|
||||
LicenseUsageServices._validate_embedding_quota(license_period, file_size_mb)
|
||||
|
||||
@staticmethod
|
||||
def check_interaction_quota(tenant_id: int) -> None:
|
||||
"""
|
||||
Check if a tenant can execute a specialist without exceeding interaction quota. As it is impossible to estimate
|
||||
the number of interaction tokens, we only check if the interaction quota are exceeded. So we might have a
|
||||
limited overrun.
|
||||
|
||||
Args:
|
||||
tenant_id: ID of the tenant
|
||||
|
||||
Raises:
|
||||
EveAIInteractionQuotaExceeded: If interaction quota would be exceeded
|
||||
EveAINoActiveLicense: If no active license is found
|
||||
EveAIException: For other errors
|
||||
"""
|
||||
# Get active license period
|
||||
license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
|
||||
# Early return if both overruns are allowed - no need to check usage at all
|
||||
if license_period.additional_interaction_allowed:
|
||||
return
|
||||
|
||||
# Convert tokens to M tokens and check interaction quota
|
||||
LicenseUsageServices._validate_interaction_quota(license_period)
|
||||
|
||||
@staticmethod
|
||||
def _validate_storage_quota(license_period: LicensePeriod, additional_mb: float) -> None:
|
||||
"""Check storage quota and raise exception if exceeded"""
|
||||
current_storage = license_period.license_usage.storage_mb_used or 0
|
||||
projected_storage = current_storage + additional_mb
|
||||
max_storage = license_period.max_storage_mb
|
||||
|
||||
# Hard limit check (we only get here if overruns are NOT allowed)
|
||||
if projected_storage > max_storage:
|
||||
raise EveAIStorageQuotaExceeded(
|
||||
current_usage=current_storage,
|
||||
limit=max_storage,
|
||||
additional=additional_mb
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_embedding_quota(license_period: LicensePeriod, additional_mb: float) -> None:
|
||||
"""Check embedding quota and raise exception if exceeded"""
|
||||
current_embedding = license_period.license_usage.embedding_mb_used or 0
|
||||
projected_embedding = current_embedding + additional_mb
|
||||
max_embedding = license_period.included_embedding_mb
|
||||
|
||||
# Hard limit check (we only get here if overruns are NOT allowed)
|
||||
if projected_embedding > max_embedding:
|
||||
raise EveAIEmbeddingQuotaExceeded(
|
||||
current_usage=current_embedding,
|
||||
limit=max_embedding,
|
||||
additional=additional_mb
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_interaction_quota(license_period) -> None:
|
||||
"""Check interaction quota and raise exception if exceeded (tokens in millions). We might have an overrun!"""
|
||||
current_tokens = license_period.license_usage.interaction_total_tokens_used / 1_000_000 or 0
|
||||
max_tokens = license_period.included_interaction_tokens
|
||||
|
||||
# Hard limit check (we only get here if overruns are NOT allowed)
|
||||
if current_tokens > max_tokens:
|
||||
raise EveAIInteractionQuotaExceeded(
|
||||
current_usage=current_tokens,
|
||||
limit=max_tokens
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
222
common/services/interaction/specialist_services.py
Normal file
222
common/services/interaction/specialist_services.py
Normal file
@@ -0,0 +1,222 @@
|
||||
import uuid
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
from typing import Dict, Any, Tuple, Optional
|
||||
from flask import current_app
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from common.extensions import db, cache_manager
|
||||
from common.models.interaction import (
|
||||
Specialist, EveAIAgent, EveAITask, EveAITool
|
||||
)
|
||||
from common.utils.celery_utils import current_celery
|
||||
from common.utils.model_logging_utils import set_logging_information, update_logging_information
|
||||
|
||||
|
||||
class SpecialistServices:
|
||||
@staticmethod
|
||||
def start_session() -> str:
|
||||
return f"CHAT_SESSION_{uuid.uuid4()}"
|
||||
|
||||
@staticmethod
|
||||
def execute_specialist(tenant_id, specialist_id, specialist_arguments, session_id, user_timezone) -> Dict[str, Any]:
|
||||
task = current_celery.send_task(
|
||||
'execute_specialist',
|
||||
args=[tenant_id,
|
||||
specialist_id,
|
||||
specialist_arguments,
|
||||
session_id,
|
||||
user_timezone,
|
||||
],
|
||||
queue='llm_interactions'
|
||||
)
|
||||
|
||||
return {
|
||||
'task_id': task.id,
|
||||
'status': 'queued',
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def initialize_specialist(specialist_id: int, specialist_type: str, specialist_version: str):
|
||||
"""
|
||||
Initialize an agentic specialist by creating all its components based on configuration.
|
||||
|
||||
Args:
|
||||
specialist_id: ID of the specialist to initialize
|
||||
specialist_type: Type of the specialist
|
||||
specialist_version: Version of the specialist type to use
|
||||
|
||||
Raises:
|
||||
ValueError: If specialist not found or invalid configuration
|
||||
SQLAlchemyError: If database operations fail
|
||||
"""
|
||||
config = cache_manager.specialists_config_cache.get_config(specialist_type, specialist_version)
|
||||
if not config:
|
||||
raise ValueError(f"No configuration found for {specialist_type} version {specialist_version}")
|
||||
if config['framework'] == 'langchain':
|
||||
pass # Langchain does not require additional items to be initialized. All configuration is in the specialist.
|
||||
|
||||
specialist = Specialist.query.get(specialist_id)
|
||||
if not specialist:
|
||||
raise ValueError(f"Specialist with ID {specialist_id} not found")
|
||||
|
||||
if config['framework'] == 'crewai':
|
||||
SpecialistServices.initialize_crewai_specialist(specialist, config)
|
||||
|
||||
@staticmethod
|
||||
def initialize_crewai_specialist(specialist: Specialist, config: Dict[str, Any]):
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
try:
|
||||
# Initialize agents
|
||||
if 'agents' in config:
|
||||
for agent_config in config['agents']:
|
||||
SpecialistServices._create_agent(
|
||||
specialist_id=specialist.id,
|
||||
agent_type=agent_config['type'],
|
||||
agent_version=agent_config['version'],
|
||||
name=agent_config.get('name'),
|
||||
description=agent_config.get('description'),
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
# Initialize tasks
|
||||
if 'tasks' in config:
|
||||
for task_config in config['tasks']:
|
||||
SpecialistServices._create_task(
|
||||
specialist_id=specialist.id,
|
||||
task_type=task_config['type'],
|
||||
task_version=task_config['version'],
|
||||
name=task_config.get('name'),
|
||||
description=task_config.get('description'),
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
# Initialize tools
|
||||
if 'tools' in config:
|
||||
for tool_config in config['tools']:
|
||||
SpecialistServices._create_tool(
|
||||
specialist_id=specialist.id,
|
||||
tool_type=tool_config['type'],
|
||||
tool_version=tool_config['version'],
|
||||
name=tool_config.get('name'),
|
||||
description=tool_config.get('description'),
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
db.session.commit()
|
||||
current_app.logger.info(f"Successfully initialized crewai specialist {specialist.id}")
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Database error initializing crewai specialist {specialist.id}: {str(e)}")
|
||||
raise
|
||||
except Exception as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error initializing crewai specialist {specialist.id}: {str(e)}")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _create_agent(
|
||||
specialist_id: int,
|
||||
agent_type: str,
|
||||
agent_version: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
timestamp: Optional[dt] = None
|
||||
) -> EveAIAgent:
|
||||
"""Create an agent with the given configuration."""
|
||||
if timestamp is None:
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
# Get agent configuration from cache
|
||||
agent_config = cache_manager.agents_config_cache.get_config(agent_type, agent_version)
|
||||
|
||||
agent = EveAIAgent(
|
||||
specialist_id=specialist_id,
|
||||
name=name or agent_config.get('name', agent_type),
|
||||
description=description or agent_config.get('metadata').get('description', ''),
|
||||
type=agent_type,
|
||||
type_version=agent_version,
|
||||
role=None,
|
||||
goal=None,
|
||||
backstory=None,
|
||||
tuning=False,
|
||||
configuration=None,
|
||||
arguments=None
|
||||
)
|
||||
|
||||
set_logging_information(agent, timestamp)
|
||||
|
||||
db.session.add(agent)
|
||||
current_app.logger.info(f"Created agent {agent.id} of type {agent_type}")
|
||||
return agent
|
||||
|
||||
@staticmethod
|
||||
def _create_task(
|
||||
specialist_id: int,
|
||||
task_type: str,
|
||||
task_version: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
timestamp: Optional[dt] = None
|
||||
) -> EveAITask:
|
||||
"""Create a task with the given configuration."""
|
||||
if timestamp is None:
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
# Get task configuration from cache
|
||||
task_config = cache_manager.tasks_config_cache.get_config(task_type, task_version)
|
||||
|
||||
task = EveAITask(
|
||||
specialist_id=specialist_id,
|
||||
name=name or task_config.get('name', task_type),
|
||||
description=description or task_config.get('metadata').get('description', ''),
|
||||
type=task_type,
|
||||
type_version=task_version,
|
||||
task_description=None,
|
||||
expected_output=None,
|
||||
tuning=False,
|
||||
configuration=None,
|
||||
arguments=None,
|
||||
context=None,
|
||||
asynchronous=False,
|
||||
)
|
||||
|
||||
set_logging_information(task, timestamp)
|
||||
|
||||
db.session.add(task)
|
||||
current_app.logger.info(f"Created task {task.id} of type {task_type}")
|
||||
return task
|
||||
|
||||
@staticmethod
|
||||
def _create_tool(
|
||||
specialist_id: int,
|
||||
tool_type: str,
|
||||
tool_version: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
timestamp: Optional[dt] = None
|
||||
) -> EveAITool:
|
||||
"""Create a tool with the given configuration."""
|
||||
if timestamp is None:
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
# Get tool configuration from cache
|
||||
tool_config = cache_manager.tools_config_cache.get_config(tool_type, tool_version)
|
||||
|
||||
tool = EveAITool(
|
||||
specialist_id=specialist_id,
|
||||
name=name or tool_config.get('name', tool_type),
|
||||
description=description or tool_config.get('metadata').get('description', ''),
|
||||
type=tool_type,
|
||||
type_version=tool_version,
|
||||
tuning=False,
|
||||
configuration=None,
|
||||
arguments=None,
|
||||
)
|
||||
|
||||
set_logging_information(tool, timestamp)
|
||||
|
||||
db.session.add(tool)
|
||||
current_app.logger.info(f"Created tool {tool.id} of type {tool_type}")
|
||||
return tool
|
||||
5
common/services/user/__init__.py
Normal file
5
common/services/user/__init__.py
Normal file
@@ -0,0 +1,5 @@
|
||||
from common.services.user.user_services import UserServices
|
||||
from common.services.user.partner_services import PartnerServices
|
||||
from common.services.user.tenant_services import TenantServices
|
||||
|
||||
__all__ = ['UserServices', 'PartnerServices', 'TenantServices']
|
||||
47
common/services/user/partner_services.py
Normal file
47
common/services/user/partner_services.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from typing import List
|
||||
|
||||
from flask import session
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from common.models.entitlements import PartnerServiceLicenseTier
|
||||
from common.utils.eveai_exceptions import EveAINoManagementPartnerService, EveAINoSessionPartner
|
||||
|
||||
from common.utils.security_utils import current_user_has_role
|
||||
|
||||
|
||||
class PartnerServices:
|
||||
@staticmethod
|
||||
def get_allowed_license_tier_ids() -> List[int]:
|
||||
"""
|
||||
Retrieve IDs of all License Tiers associated with the partner's management service
|
||||
|
||||
Returns:
|
||||
List of license tier IDs
|
||||
|
||||
Raises:
|
||||
EveAINoSessionPartner: If no partner is in the session
|
||||
EveAINoManagementPartnerService: If partner has no management service
|
||||
"""
|
||||
partner = session.get("partner", None)
|
||||
if not partner:
|
||||
raise EveAINoSessionPartner()
|
||||
|
||||
# Find a management service for this partner
|
||||
management_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
if not management_service:
|
||||
raise EveAINoManagementPartnerService()
|
||||
management_service_id = management_service['id']
|
||||
|
||||
# Query for all license tiers associated with this management service
|
||||
associations = PartnerServiceLicenseTier.query.filter_by(
|
||||
partner_service_id=management_service_id
|
||||
).all()
|
||||
|
||||
# Extract the license tier IDs
|
||||
license_tier_ids = [assoc.license_tier_id for assoc in associations]
|
||||
|
||||
return license_tier_ids
|
||||
|
||||
|
||||
|
||||
175
common/services/user/tenant_services.py
Normal file
175
common/services/user/tenant_services.py
Normal file
@@ -0,0 +1,175 @@
|
||||
from typing import Dict, List
|
||||
|
||||
from flask import session, current_app
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from common.extensions import db, cache_manager
|
||||
from common.models.user import Partner, PartnerTenant, PartnerService, Tenant
|
||||
from common.utils.eveai_exceptions import EveAINoManagementPartnerService
|
||||
from common.utils.model_logging_utils import set_logging_information
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
|
||||
from common.utils.security_utils import current_user_has_role
|
||||
|
||||
|
||||
class TenantServices:
|
||||
@staticmethod
|
||||
def associate_tenant_with_partner(tenant_id):
|
||||
"""Associate a tenant with a partner"""
|
||||
try:
|
||||
partner_id = session['partner']['id']
|
||||
# Get partner service (MANAGEMENT_SERVICE type)
|
||||
partner = Partner.query.get(partner_id)
|
||||
if not partner:
|
||||
return
|
||||
|
||||
# Find a management service for this partner
|
||||
management_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
|
||||
if not management_service:
|
||||
current_app.logger.error(f"No Management Service defined for partner {partner_id}"
|
||||
f"while associating tenant {tenant_id} with partner.")
|
||||
raise EveAINoManagementPartnerService()
|
||||
|
||||
# Create the association
|
||||
tenant_partner = PartnerTenant(
|
||||
partner_service_id=management_service['id'],
|
||||
tenant_id=tenant_id,
|
||||
)
|
||||
set_logging_information(tenant_partner, dt.now(tz.utc))
|
||||
|
||||
db.session.add(tenant_partner)
|
||||
db.session.commit()
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error associating tenant {tenant_id} with partner: {str(e)}")
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def get_available_types_for_tenant(tenant_id: int, config_type: str) -> Dict[str, Dict[str, str]]:
|
||||
"""
|
||||
Get available configuration types for a tenant based on partner relationships
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant ID
|
||||
config_type: The configuration type ('specialists', 'agents', 'tasks', etc.)
|
||||
|
||||
Returns:
|
||||
Dictionary of available types for the tenant
|
||||
"""
|
||||
# Get the appropriate cache handler based on config_type
|
||||
cache_handler = None
|
||||
if config_type == 'specialists':
|
||||
cache_handler = cache_manager.specialists_types_cache
|
||||
elif config_type == 'agents':
|
||||
cache_handler = cache_manager.agents_types_cache
|
||||
elif config_type == 'tasks':
|
||||
cache_handler = cache_manager.tasks_types_cache
|
||||
elif config_type == 'tools':
|
||||
cache_handler = cache_manager.tools_types_cache
|
||||
else:
|
||||
raise ValueError(f"Unsupported config type: {config_type}")
|
||||
|
||||
# Get all types with their metadata (including partner info)
|
||||
all_types = cache_handler.get_types()
|
||||
|
||||
# Filter to include:
|
||||
# 1. Types with no partner (global)
|
||||
# 2. Types with partners that have a SPECIALIST_SERVICE relationship with this tenant
|
||||
available_partners = TenantServices.get_tenant_partner_names(tenant_id)
|
||||
|
||||
available_types = {
|
||||
type_id: info for type_id, info in all_types.items()
|
||||
if info.get('partner') is None or info.get('partner') in available_partners
|
||||
}
|
||||
|
||||
return available_types
|
||||
|
||||
@staticmethod
|
||||
def get_tenant_partner_names(tenant_id: int) -> List[str]:
|
||||
"""
|
||||
Get names of partners that have a SPECIALIST_SERVICE relationship with this tenant
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant ID
|
||||
|
||||
Returns:
|
||||
List of partner names (tenant names)
|
||||
"""
|
||||
# Find all PartnerTenant relationships for this tenant
|
||||
partner_names = []
|
||||
try:
|
||||
# Get all partner services of type SPECIALIST_SERVICE
|
||||
specialist_services = (
|
||||
PartnerService.query
|
||||
.filter_by(type='SPECIALIST_SERVICE')
|
||||
.all()
|
||||
)
|
||||
|
||||
if not specialist_services:
|
||||
return []
|
||||
|
||||
# Find tenant relationships with these services
|
||||
partner_tenants = (
|
||||
PartnerTenant.query
|
||||
.filter_by(tenant_id=tenant_id)
|
||||
.filter(PartnerTenant.partner_service_id.in_([svc.id for svc in specialist_services]))
|
||||
.all()
|
||||
)
|
||||
|
||||
# Get the partner names (their tenant names)
|
||||
for pt in partner_tenants:
|
||||
partner_service = (
|
||||
PartnerService.query
|
||||
.filter_by(id=pt.partner_service_id)
|
||||
.first()
|
||||
)
|
||||
|
||||
if partner_service:
|
||||
partner = Partner.query.get(partner_service.partner_id)
|
||||
if partner:
|
||||
# Get the tenant associated with this partner
|
||||
partner_tenant = Tenant.query.get(partner.tenant_id)
|
||||
if partner_tenant:
|
||||
partner_names.append(partner_tenant.name)
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f"Database error retrieving partner names: {str(e)}")
|
||||
|
||||
return partner_names
|
||||
|
||||
@staticmethod
|
||||
def can_use_specialist_type(tenant_id: int, specialist_type: str) -> bool:
|
||||
"""
|
||||
Check if a tenant can use a specific specialist type
|
||||
|
||||
Args:
|
||||
tenant_id: The tenant ID
|
||||
specialist_type: The specialist type ID
|
||||
|
||||
Returns:
|
||||
True if the tenant can use the specialist type, False otherwise
|
||||
"""
|
||||
# Get the specialist type definition
|
||||
try:
|
||||
specialist_types = cache_manager.specialists_types_cache.get_types()
|
||||
specialist_def = specialist_types.get(specialist_type)
|
||||
|
||||
if not specialist_def:
|
||||
return False
|
||||
|
||||
# If it's a global specialist, anyone can use it
|
||||
if specialist_def.get('partner') is None:
|
||||
return True
|
||||
|
||||
# If it's a partner-specific specialist, check if tenant has access
|
||||
partner_name = specialist_def.get('partner')
|
||||
available_partners = TenantServices.get_tenant_partner_names(tenant_id)
|
||||
|
||||
return partner_name in available_partners
|
||||
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error checking specialist type access: {str(e)}")
|
||||
return False
|
||||
95
common/services/user/user_services.py
Normal file
95
common/services/user/user_services.py
Normal file
@@ -0,0 +1,95 @@
|
||||
from flask import session
|
||||
|
||||
from common.models.user import Partner, Role, PartnerTenant
|
||||
|
||||
from common.utils.eveai_exceptions import EveAIRoleAssignmentException
|
||||
from common.utils.security_utils import current_user_has_role
|
||||
|
||||
|
||||
class UserServices:
|
||||
@staticmethod
|
||||
def get_assignable_roles():
|
||||
"""Retrieves roles that can be assigned to a user depending on the current user logged in,
|
||||
and the active tenant for the session"""
|
||||
current_tenant_id = session.get('tenant').get('id', None)
|
||||
effective_role_names = []
|
||||
if current_tenant_id == 1:
|
||||
if current_user_has_role("Super User"):
|
||||
effective_role_names.append("Super User")
|
||||
elif current_tenant_id:
|
||||
if current_user_has_role("Tenant Admin"):
|
||||
effective_role_names.append("Tenant Admin")
|
||||
if current_user_has_role("Partner Admin") or current_user_has_role("Super User"):
|
||||
effective_role_names.append("Tenant Admin")
|
||||
if session.get('partner'):
|
||||
if session.get('partner').get('tenant_id') == current_tenant_id:
|
||||
effective_role_names.append("Partner Admin")
|
||||
effective_role_names = list(set(effective_role_names))
|
||||
effective_roles = [(role.id, role.name) for role in
|
||||
Role.query.filter(Role.name.in_(effective_role_names)).all()]
|
||||
return effective_roles
|
||||
|
||||
@staticmethod
|
||||
def validate_role_assignments(role_ids):
|
||||
"""Validate a set of role assignments, raising exception for first invalid role"""
|
||||
assignable_roles = UserServices.get_assignable_roles()
|
||||
assignable_role_ids = {role[0] for role in assignable_roles}
|
||||
role_id_set = set(role_ids)
|
||||
return role_id_set.issubset(assignable_role_ids)
|
||||
|
||||
@staticmethod
|
||||
def can_user_edit_tenant(tenant_id) -> bool:
|
||||
if current_user_has_role('Super User'):
|
||||
return True
|
||||
elif current_user_has_role('Partner Admin'):
|
||||
partner = session.get('partner', None)
|
||||
if partner and partner["tenant_id"] == tenant_id:
|
||||
return True
|
||||
partner_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
if not partner_service:
|
||||
return False
|
||||
else:
|
||||
partner_tenant = PartnerTenant.query.filter(
|
||||
PartnerTenant.tenant_id == tenant_id,
|
||||
PartnerTenant.partner_service_id == partner_service['id'],
|
||||
).first()
|
||||
if partner_tenant:
|
||||
return True
|
||||
else:
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def can_user_create_tenant() -> bool:
|
||||
if current_user_has_role('Super User'):
|
||||
return True
|
||||
elif current_user_has_role('Partner Admin'):
|
||||
partner_id = session['partner']['id']
|
||||
partner_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
if not partner_service:
|
||||
return False
|
||||
else:
|
||||
partner_permissions = partner_service.get('permissions', None)
|
||||
return partner_permissions.get('can_create_tenant', False)
|
||||
else:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def can_user_assign_license() -> bool:
|
||||
if current_user_has_role('Super User'):
|
||||
return True
|
||||
elif current_user_has_role('Partner Admin'):
|
||||
partner_id = session['partner']['id']
|
||||
partner_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
if not partner_service:
|
||||
return False
|
||||
else:
|
||||
partner_permissions = partner_service.get('permissions', None)
|
||||
return partner_permissions.get('can_assign_license', False)
|
||||
else:
|
||||
return False
|
||||
|
||||
BIN
common/utils/.DS_Store
vendored
BIN
common/utils/.DS_Store
vendored
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
62
common/utils/asset_utils.py
Normal file
62
common/utils/asset_utils.py
Normal file
@@ -0,0 +1,62 @@
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
|
||||
from flask import current_app
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from common.extensions import cache_manager, minio_client, db
|
||||
from common.models.interaction import EveAIAsset, EveAIAssetVersion
|
||||
from common.utils.model_logging_utils import set_logging_information
|
||||
|
||||
|
||||
def create_asset_stack(api_input, tenant_id):
|
||||
type_version = cache_manager.assets_version_tree_cache.get_latest_version(api_input['type'])
|
||||
api_input['type_version'] = type_version
|
||||
new_asset = create_asset(api_input, tenant_id)
|
||||
new_asset_version = create_version_for_asset(new_asset, tenant_id)
|
||||
db.session.add(new_asset)
|
||||
db.session.add(new_asset_version)
|
||||
|
||||
try:
|
||||
db.session.commit()
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f"Could not add asset for tenant {tenant_id}: {str(e)}")
|
||||
db.session.rollback()
|
||||
raise e
|
||||
|
||||
return new_asset, new_asset_version
|
||||
|
||||
|
||||
def create_asset(api_input, tenant_id):
|
||||
new_asset = EveAIAsset()
|
||||
new_asset.name = api_input['name']
|
||||
new_asset.description = api_input['description']
|
||||
new_asset.type = api_input['type']
|
||||
new_asset.type_version = api_input['type_version']
|
||||
if api_input['valid_from'] and api_input['valid_from'] != '':
|
||||
new_asset.valid_from = api_input['valid_from']
|
||||
else:
|
||||
new_asset.valid_from = dt.now(tz.utc)
|
||||
new_asset.valid_to = api_input['valid_to']
|
||||
set_logging_information(new_asset, dt.now(tz.utc))
|
||||
|
||||
return new_asset
|
||||
|
||||
|
||||
def create_version_for_asset(asset, tenant_id):
|
||||
new_asset_version = EveAIAssetVersion()
|
||||
new_asset_version.asset = asset
|
||||
new_asset_version.bucket_name = minio_client.create_tenant_bucket(tenant_id)
|
||||
set_logging_information(new_asset_version, dt.now(tz.utc))
|
||||
|
||||
return new_asset_version
|
||||
|
||||
|
||||
def add_asset_version_file(asset_version, field_name, file, tenant_id):
|
||||
object_name, file_size = minio_client.upload_file(asset_version.bucket_name, asset_version.id, field_name,
|
||||
file.content_type)
|
||||
# mark_tenant_storage_dirty(tenant_id)
|
||||
# TODO - zorg ervoor dat de herberekening van storage onmiddellijk gebeurt!
|
||||
return object_name
|
||||
|
||||
|
||||
|
||||
@@ -1,15 +1,83 @@
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from contextlib import contextmanager, asynccontextmanager
|
||||
from datetime import datetime
|
||||
from typing import Dict, Any, Optional
|
||||
from typing import Dict, Any, Optional, List
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
from portkey_ai import Portkey, Config
|
||||
import logging
|
||||
|
||||
from flask import current_app
|
||||
from prometheus_client import Counter, Histogram, Gauge, Summary, push_to_gateway, REGISTRY
|
||||
|
||||
from .business_event_context import BusinessEventContext
|
||||
from common.models.entitlements import BusinessEventLog
|
||||
from common.extensions import db
|
||||
from .celery_utils import current_celery
|
||||
from common.utils.prometheus_utils import sanitize_label
|
||||
|
||||
# Standard duration buckets for all histograms
|
||||
DURATION_BUCKETS = [0.1, 0.5, 1, 2.5, 5, 10, 15, 30, 60, 120, 240, 360, float('inf')]
|
||||
|
||||
# Prometheus metrics for business events
|
||||
TRACE_COUNTER = Counter(
|
||||
'eveai_business_events_total',
|
||||
'Total number of business events triggered',
|
||||
['tenant_id', 'event_type', 'specialist_id', 'specialist_type', 'specialist_type_version']
|
||||
)
|
||||
|
||||
TRACE_DURATION = Histogram(
|
||||
'eveai_business_events_duration_seconds',
|
||||
'Duration of business events in seconds',
|
||||
['tenant_id', 'event_type', 'specialist_id', 'specialist_type', 'specialist_type_version'],
|
||||
buckets=DURATION_BUCKETS
|
||||
)
|
||||
|
||||
CONCURRENT_TRACES = Gauge(
|
||||
'eveai_business_events_concurrent',
|
||||
'Number of concurrent business events',
|
||||
['tenant_id', 'event_type', 'specialist_id', 'specialist_type', 'specialist_type_version']
|
||||
)
|
||||
|
||||
SPAN_COUNTER = Counter(
|
||||
'eveai_business_spans_total',
|
||||
'Total number of spans within business events',
|
||||
['tenant_id', 'event_type', 'activity_name', 'specialist_id', 'specialist_type', 'specialist_type_version']
|
||||
)
|
||||
|
||||
SPAN_DURATION = Histogram(
|
||||
'eveai_business_spans_duration_seconds',
|
||||
'Duration of spans within business events in seconds',
|
||||
['tenant_id', 'event_type', 'activity_name', 'specialist_id', 'specialist_type', 'specialist_type_version'],
|
||||
buckets=DURATION_BUCKETS
|
||||
)
|
||||
|
||||
CONCURRENT_SPANS = Gauge(
|
||||
'eveai_business_spans_concurrent',
|
||||
'Number of concurrent spans within business events',
|
||||
['tenant_id', 'event_type', 'activity_name', 'specialist_id', 'specialist_type', 'specialist_type_version']
|
||||
)
|
||||
|
||||
# LLM Usage metrics
|
||||
LLM_TOKENS_COUNTER = Counter(
|
||||
'eveai_llm_tokens_total',
|
||||
'Total number of tokens used in LLM calls',
|
||||
['tenant_id', 'event_type', 'interaction_type', 'token_type', 'specialist_id', 'specialist_type',
|
||||
'specialist_type_version']
|
||||
)
|
||||
|
||||
LLM_DURATION = Histogram(
|
||||
'eveai_llm_duration_seconds',
|
||||
'Duration of LLM API calls in seconds',
|
||||
['tenant_id', 'event_type', 'interaction_type', 'specialist_id', 'specialist_type', 'specialist_type_version'],
|
||||
buckets=DURATION_BUCKETS
|
||||
)
|
||||
|
||||
LLM_CALLS_COUNTER = Counter(
|
||||
'eveai_llm_calls_total',
|
||||
'Total number of LLM API calls',
|
||||
['tenant_id', 'event_type', 'interaction_type', 'specialist_id', 'specialist_type', 'specialist_type_version']
|
||||
)
|
||||
|
||||
|
||||
class BusinessEvent:
|
||||
@@ -28,6 +96,9 @@ class BusinessEvent:
|
||||
self.document_version_file_size = kwargs.get('document_version_file_size')
|
||||
self.chat_session_id = kwargs.get('chat_session_id')
|
||||
self.interaction_id = kwargs.get('interaction_id')
|
||||
self.specialist_id = kwargs.get('specialist_id')
|
||||
self.specialist_type = kwargs.get('specialist_type')
|
||||
self.specialist_type_version = kwargs.get('specialist_type_version')
|
||||
self.environment = os.environ.get("FLASK_ENV", "development")
|
||||
self.span_counter = 0
|
||||
self.spans = []
|
||||
@@ -35,29 +106,131 @@ class BusinessEvent:
|
||||
'total_tokens': 0,
|
||||
'prompt_tokens': 0,
|
||||
'completion_tokens': 0,
|
||||
'nr_of_pages': 0,
|
||||
'total_time': 0,
|
||||
'call_count': 0,
|
||||
'interaction_type': None
|
||||
}
|
||||
self._log_buffer = []
|
||||
|
||||
# Prometheus label values must be strings
|
||||
self.tenant_id_str = str(self.tenant_id)
|
||||
self.event_type_str = sanitize_label(self.event_type)
|
||||
self.specialist_id_str = str(self.specialist_id) if self.specialist_id else ""
|
||||
self.specialist_type_str = str(self.specialist_type) if self.specialist_type else ""
|
||||
self.specialist_type_version_str = sanitize_label(str(self.specialist_type_version)) \
|
||||
if self.specialist_type_version else ""
|
||||
self.span_name_str = ""
|
||||
|
||||
# Increment concurrent events gauge when initialized
|
||||
CONCURRENT_TRACES.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
# Increment trace counter
|
||||
TRACE_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
def update_attribute(self, attribute: str, value: any):
|
||||
if hasattr(self, attribute):
|
||||
setattr(self, attribute, value)
|
||||
# Update string versions for Prometheus labels if needed
|
||||
if attribute == 'specialist_id':
|
||||
self.specialist_id_str = str(value) if value else ""
|
||||
elif attribute == 'specialist_type':
|
||||
self.specialist_type_str = str(value) if value else ""
|
||||
elif attribute == 'specialist_type_version':
|
||||
self.specialist_type_version_str = sanitize_label(str(value)) if value else ""
|
||||
elif attribute == 'tenant_id':
|
||||
self.tenant_id_str = str(value)
|
||||
elif attribute == 'event_type':
|
||||
self.event_type_str = sanitize_label(value)
|
||||
elif attribute == 'span_name':
|
||||
self.span_name_str = sanitize_label(value)
|
||||
else:
|
||||
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attribute}'")
|
||||
|
||||
def update_llm_metrics(self, metrics: dict):
|
||||
self.llm_metrics['total_tokens'] += metrics['total_tokens']
|
||||
self.llm_metrics['prompt_tokens'] += metrics['prompt_tokens']
|
||||
self.llm_metrics['completion_tokens'] += metrics['completion_tokens']
|
||||
self.llm_metrics['total_time'] += metrics['time_elapsed']
|
||||
self.llm_metrics['total_tokens'] += metrics.get('total_tokens', 0)
|
||||
self.llm_metrics['prompt_tokens'] += metrics.get('prompt_tokens', 0)
|
||||
self.llm_metrics['completion_tokens'] += metrics.get('completion_tokens', 0)
|
||||
self.llm_metrics['nr_of_pages'] += metrics.get('nr_of_pages', 0)
|
||||
self.llm_metrics['total_time'] += metrics.get('time_elapsed', 0)
|
||||
self.llm_metrics['call_count'] += 1
|
||||
self.llm_metrics['interaction_type'] = metrics['interaction_type']
|
||||
|
||||
# Track in Prometheus metrics
|
||||
interaction_type_str = sanitize_label(metrics['interaction_type']) if metrics['interaction_type'] else ""
|
||||
|
||||
# Track token usage
|
||||
LLM_TOKENS_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
interaction_type=interaction_type_str,
|
||||
token_type='total',
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc(metrics.get('total_tokens', 0))
|
||||
|
||||
LLM_TOKENS_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
interaction_type=interaction_type_str,
|
||||
token_type='prompt',
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc(metrics.get('prompt_tokens', 0))
|
||||
|
||||
LLM_TOKENS_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
interaction_type=interaction_type_str,
|
||||
token_type='completion',
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc(metrics.get('completion_tokens', 0))
|
||||
|
||||
# Track duration
|
||||
LLM_DURATION.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
interaction_type=interaction_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).observe(metrics.get('time_elapsed', 0))
|
||||
|
||||
# Track call count
|
||||
LLM_CALLS_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
interaction_type=interaction_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
def reset_llm_metrics(self):
|
||||
self.llm_metrics['total_tokens'] = 0
|
||||
self.llm_metrics['prompt_tokens'] = 0
|
||||
self.llm_metrics['completion_tokens'] = 0
|
||||
self.llm_metrics['nr_of_pages'] = 0
|
||||
self.llm_metrics['total_time'] = 0
|
||||
self.llm_metrics['call_count'] = 0
|
||||
self.llm_metrics['interaction_type'] = None
|
||||
@@ -79,28 +252,166 @@ class BusinessEvent:
|
||||
# Set the new span info
|
||||
self.span_id = new_span_id
|
||||
self.span_name = span_name
|
||||
self.span_name_str = sanitize_label(span_name) if span_name else ""
|
||||
self.parent_span_id = parent_span_id
|
||||
|
||||
self.log(f"Starting span {span_name}")
|
||||
# Track start time for the span
|
||||
span_start_time = time.time()
|
||||
|
||||
# Increment span metrics - using span_name as activity_name for metrics
|
||||
SPAN_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
# Increment concurrent spans gauge
|
||||
CONCURRENT_SPANS.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
self.log(f"Start")
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Calculate total time for this span
|
||||
span_total_time = time.time() - span_start_time
|
||||
|
||||
# Observe span duration
|
||||
SPAN_DURATION.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).observe(span_total_time)
|
||||
|
||||
# Decrement concurrent spans gauge
|
||||
CONCURRENT_SPANS.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).dec()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
if self.llm_metrics['call_count'] > 0:
|
||||
self.log_final_metrics()
|
||||
self.reset_llm_metrics()
|
||||
self.log(f"Ending span {span_name}")
|
||||
self.log(f"End", extra_fields={'span_duration': span_total_time})
|
||||
# Restore the previous span info
|
||||
if self.spans:
|
||||
self.span_id, self.span_name, self.parent_span_id = self.spans.pop()
|
||||
self.span_name_str = sanitize_label(span_name) if span_name else ""
|
||||
else:
|
||||
self.span_id = None
|
||||
self.span_name = None
|
||||
self.parent_span_id = None
|
||||
self.span_name_str = ""
|
||||
|
||||
def log(self, message: str, level: str = 'info'):
|
||||
logger = logging.getLogger('business_events')
|
||||
@asynccontextmanager
|
||||
async def create_span_async(self, span_name: str):
|
||||
"""Async version of create_span using async context manager"""
|
||||
parent_span_id = self.span_id
|
||||
self.span_counter += 1
|
||||
new_span_id = str(uuid.uuid4())
|
||||
|
||||
# Save the current span info
|
||||
self.spans.append((self.span_id, self.span_name, self.parent_span_id))
|
||||
|
||||
# Set the new span info
|
||||
self.span_id = new_span_id
|
||||
self.span_name = span_name
|
||||
self.span_name_str = sanitize_label(span_name) if span_name else ""
|
||||
self.parent_span_id = parent_span_id
|
||||
|
||||
# Track start time for the span
|
||||
span_start_time = time.time()
|
||||
|
||||
# Increment span metrics - using span_name as activity_name for metrics
|
||||
SPAN_COUNTER.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
# Increment concurrent spans gauge
|
||||
CONCURRENT_SPANS.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).inc()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
self.log(f"Start")
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Calculate total time for this span
|
||||
span_total_time = time.time() - span_start_time
|
||||
|
||||
# Observe span duration
|
||||
SPAN_DURATION.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).observe(span_total_time)
|
||||
|
||||
# Decrement concurrent spans gauge
|
||||
CONCURRENT_SPANS.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
activity_name=self.span_name_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).dec()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
if self.llm_metrics['call_count'] > 0:
|
||||
self.log_final_metrics()
|
||||
self.reset_llm_metrics()
|
||||
self.log(f"End", extra_fields={'span_duration': span_total_time})
|
||||
# Restore the previous span info
|
||||
if self.spans:
|
||||
self.span_id, self.span_name, self.parent_span_id = self.spans.pop()
|
||||
self.span_name_str = sanitize_label(span_name) if span_name else ""
|
||||
else:
|
||||
self.span_id = None
|
||||
self.span_name = None
|
||||
self.parent_span_id = None
|
||||
self.span_name_str = ""
|
||||
|
||||
def log(self, message: str, level: str = 'info', extra_fields: Dict[str, Any] = None):
|
||||
log_data = {
|
||||
'timestamp': dt.now(tz=tz.utc),
|
||||
'event_type': self.event_type,
|
||||
'tenant_id': self.tenant_id,
|
||||
'trace_id': self.trace_id,
|
||||
@@ -111,35 +422,29 @@ class BusinessEvent:
|
||||
'document_version_file_size': self.document_version_file_size,
|
||||
'chat_session_id': self.chat_session_id,
|
||||
'interaction_id': self.interaction_id,
|
||||
'specialist_id': self.specialist_id,
|
||||
'specialist_type': self.specialist_type,
|
||||
'specialist_type_version': self.specialist_type_version,
|
||||
'environment': self.environment,
|
||||
'message': message,
|
||||
}
|
||||
# log to Graylog
|
||||
getattr(logger, level)(message, extra=log_data)
|
||||
# Add any extra fields
|
||||
if extra_fields:
|
||||
for key, value in extra_fields.items():
|
||||
# For span/trace duration, use the llm_metrics_total_time field
|
||||
if key == 'span_duration' or key == 'trace_duration':
|
||||
log_data['llm_metrics_total_time'] = value
|
||||
else:
|
||||
log_data[key] = value
|
||||
|
||||
# Log to database
|
||||
event_log = BusinessEventLog(
|
||||
timestamp=dt.now(tz=tz.utc),
|
||||
event_type=self.event_type,
|
||||
tenant_id=self.tenant_id,
|
||||
trace_id=self.trace_id,
|
||||
span_id=self.span_id,
|
||||
span_name=self.span_name,
|
||||
parent_span_id=self.parent_span_id,
|
||||
document_version_id=self.document_version_id,
|
||||
document_version_file_size=self.document_version_file_size,
|
||||
chat_session_id=self.chat_session_id,
|
||||
interaction_id=self.interaction_id,
|
||||
environment=self.environment,
|
||||
message=message
|
||||
)
|
||||
db.session.add(event_log)
|
||||
db.session.commit()
|
||||
self._log_buffer.append(log_data)
|
||||
|
||||
def log_llm_metrics(self, metrics: dict, level: str = 'info'):
|
||||
self.update_llm_metrics(metrics)
|
||||
message = "LLM Metrics"
|
||||
logger = logging.getLogger('business_events')
|
||||
log_data = {
|
||||
'timestamp': dt.now(tz=tz.utc),
|
||||
'event_type': self.event_type,
|
||||
'tenant_id': self.tenant_id,
|
||||
'trace_id': self.trace_id,
|
||||
@@ -150,44 +455,25 @@ class BusinessEvent:
|
||||
'document_version_file_size': self.document_version_file_size,
|
||||
'chat_session_id': self.chat_session_id,
|
||||
'interaction_id': self.interaction_id,
|
||||
'specialist_id': self.specialist_id,
|
||||
'specialist_type': self.specialist_type,
|
||||
'specialist_type_version': self.specialist_type_version,
|
||||
'environment': self.environment,
|
||||
'llm_metrics_total_tokens': metrics['total_tokens'],
|
||||
'llm_metrics_prompt_tokens': metrics['prompt_tokens'],
|
||||
'llm_metrics_completion_tokens': metrics['completion_tokens'],
|
||||
'llm_metrics_total_time': metrics['time_elapsed'],
|
||||
'llm_metrics_total_tokens': metrics.get('total_tokens', 0),
|
||||
'llm_metrics_prompt_tokens': metrics.get('prompt_tokens', 0),
|
||||
'llm_metrics_completion_tokens': metrics.get('completion_tokens', 0),
|
||||
'llm_metrics_nr_of_pages': metrics.get('nr_of_pages', 0),
|
||||
'llm_metrics_total_time': metrics.get('time_elapsed', 0),
|
||||
'llm_interaction_type': metrics['interaction_type'],
|
||||
'message': message,
|
||||
}
|
||||
# log to Graylog
|
||||
getattr(logger, level)(message, extra=log_data)
|
||||
|
||||
# Log to database
|
||||
event_log = BusinessEventLog(
|
||||
timestamp=dt.now(tz=tz.utc),
|
||||
event_type=self.event_type,
|
||||
tenant_id=self.tenant_id,
|
||||
trace_id=self.trace_id,
|
||||
span_id=self.span_id,
|
||||
span_name=self.span_name,
|
||||
parent_span_id=self.parent_span_id,
|
||||
document_version_id=self.document_version_id,
|
||||
document_version_file_size=self.document_version_file_size,
|
||||
chat_session_id=self.chat_session_id,
|
||||
interaction_id=self.interaction_id,
|
||||
environment=self.environment,
|
||||
llm_metrics_total_tokens=metrics['total_tokens'],
|
||||
llm_metrics_prompt_tokens=metrics['prompt_tokens'],
|
||||
llm_metrics_completion_tokens=metrics['completion_tokens'],
|
||||
llm_metrics_total_time=metrics['time_elapsed'],
|
||||
llm_interaction_type=metrics['interaction_type'],
|
||||
message=message
|
||||
)
|
||||
db.session.add(event_log)
|
||||
db.session.commit()
|
||||
self._log_buffer.append(log_data)
|
||||
|
||||
def log_final_metrics(self, level: str = 'info'):
|
||||
logger = logging.getLogger('business_events')
|
||||
message = "Final LLM Metrics"
|
||||
log_data = {
|
||||
'timestamp': dt.now(tz=tz.utc),
|
||||
'event_type': self.event_type,
|
||||
'tenant_id': self.tenant_id,
|
||||
'trace_id': self.trace_id,
|
||||
@@ -198,49 +484,161 @@ class BusinessEvent:
|
||||
'document_version_file_size': self.document_version_file_size,
|
||||
'chat_session_id': self.chat_session_id,
|
||||
'interaction_id': self.interaction_id,
|
||||
'specialist_id': self.specialist_id,
|
||||
'specialist_type': self.specialist_type,
|
||||
'specialist_type_version': self.specialist_type_version,
|
||||
'environment': self.environment,
|
||||
'llm_metrics_total_tokens': self.llm_metrics['total_tokens'],
|
||||
'llm_metrics_prompt_tokens': self.llm_metrics['prompt_tokens'],
|
||||
'llm_metrics_completion_tokens': self.llm_metrics['completion_tokens'],
|
||||
'llm_metrics_nr_of_pages': self.llm_metrics['nr_of_pages'],
|
||||
'llm_metrics_total_time': self.llm_metrics['total_time'],
|
||||
'llm_metrics_call_count': self.llm_metrics['call_count'],
|
||||
'llm_interaction_type': self.llm_metrics['interaction_type'],
|
||||
'message': message,
|
||||
}
|
||||
# log to Graylog
|
||||
getattr(logger, level)(message, extra=log_data)
|
||||
self._log_buffer.append(log_data)
|
||||
|
||||
# Log to database
|
||||
event_log = BusinessEventLog(
|
||||
timestamp=dt.now(tz=tz.utc),
|
||||
event_type=self.event_type,
|
||||
tenant_id=self.tenant_id,
|
||||
trace_id=self.trace_id,
|
||||
span_id=self.span_id,
|
||||
span_name=self.span_name,
|
||||
parent_span_id=self.parent_span_id,
|
||||
document_version_id=self.document_version_id,
|
||||
document_version_file_size=self.document_version_file_size,
|
||||
chat_session_id=self.chat_session_id,
|
||||
interaction_id=self.interaction_id,
|
||||
environment=self.environment,
|
||||
llm_metrics_total_tokens=self.llm_metrics['total_tokens'],
|
||||
llm_metrics_prompt_tokens=self.llm_metrics['prompt_tokens'],
|
||||
llm_metrics_completion_tokens=self.llm_metrics['completion_tokens'],
|
||||
llm_metrics_total_time=self.llm_metrics['total_time'],
|
||||
llm_metrics_call_count=self.llm_metrics['call_count'],
|
||||
llm_interaction_type=self.llm_metrics['interaction_type'],
|
||||
message=message
|
||||
)
|
||||
db.session.add(event_log)
|
||||
db.session.commit()
|
||||
@staticmethod
|
||||
def _direct_db_persist(log_entries: List[Dict[str, Any]]):
|
||||
"""Fallback method to directly persist logs to DB if async fails"""
|
||||
try:
|
||||
db_entries = []
|
||||
for entry in log_entries:
|
||||
event_log = BusinessEventLog(
|
||||
timestamp=entry.pop('timestamp'),
|
||||
event_type=entry.pop('event_type'),
|
||||
tenant_id=entry.pop('tenant_id'),
|
||||
trace_id=entry.pop('trace_id'),
|
||||
span_id=entry.pop('span_id', None),
|
||||
span_name=entry.pop('span_name', None),
|
||||
parent_span_id=entry.pop('parent_span_id', None),
|
||||
document_version_id=entry.pop('document_version_id', None),
|
||||
document_version_file_size=entry.pop('document_version_file_size', None),
|
||||
chat_session_id=entry.pop('chat_session_id', None),
|
||||
interaction_id=entry.pop('interaction_id', None),
|
||||
specialist_id=entry.pop('specialist_id', None),
|
||||
specialist_type=entry.pop('specialist_type', None),
|
||||
specialist_type_version=entry.pop('specialist_type_version', None),
|
||||
environment=entry.pop('environment', None),
|
||||
llm_metrics_total_tokens=entry.pop('llm_metrics_total_tokens', None),
|
||||
llm_metrics_prompt_tokens=entry.pop('llm_metrics_prompt_tokens', None),
|
||||
llm_metrics_completion_tokens=entry.pop('llm_metrics_completion_tokens', None),
|
||||
llm_metrics_total_time=entry.pop('llm_metrics_total_time', None),
|
||||
llm_metrics_call_count=entry.pop('llm_metrics_call_count', None),
|
||||
llm_interaction_type=entry.pop('llm_interaction_type', None),
|
||||
message=entry.pop('message', None)
|
||||
)
|
||||
db_entries.append(event_log)
|
||||
|
||||
# Bulk insert
|
||||
db.session.bulk_save_objects(db_entries)
|
||||
db.session.commit()
|
||||
except Exception as e:
|
||||
logger = logging.getLogger('business_events')
|
||||
logger.error(f"Failed to persist logs directly to DB: {e}")
|
||||
db.session.rollback()
|
||||
|
||||
def _flush_log_buffer(self):
|
||||
"""Flush the log buffer to the database via a Celery task"""
|
||||
if self._log_buffer:
|
||||
try:
|
||||
# Send to Celery task
|
||||
current_celery.send_task(
|
||||
'persist_business_events',
|
||||
args=[self._log_buffer],
|
||||
queue='entitlements' # Or dedicated log queue
|
||||
)
|
||||
except Exception as e:
|
||||
# Fallback to direct DB write in case of issues with Celery
|
||||
logger = logging.getLogger('business_events')
|
||||
logger.error(f"Failed to send logs to Celery. Falling back to direct DB: {e}")
|
||||
self._direct_db_persist(self._log_buffer)
|
||||
|
||||
# Clear the buffer after sending
|
||||
self._log_buffer = []
|
||||
|
||||
def _push_to_gateway(self):
|
||||
# Push metrics to the gateway
|
||||
try:
|
||||
push_to_gateway(
|
||||
current_app.config['PUSH_GATEWAY_URL'],
|
||||
job=current_app.config['COMPONENT_NAME'],
|
||||
registry=REGISTRY
|
||||
)
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Failed to push metrics to Prometheus Push Gateway: {e}")
|
||||
|
||||
def __enter__(self):
|
||||
self.trace_start_time = time.time()
|
||||
self.log(f'Starting Trace for {self.event_type}')
|
||||
return BusinessEventContext(self).__enter__()
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
trace_total_time = time.time() - self.trace_start_time
|
||||
|
||||
# Record trace duration
|
||||
TRACE_DURATION.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).observe(trace_total_time)
|
||||
|
||||
# Decrement concurrent traces gauge
|
||||
CONCURRENT_TRACES.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).dec()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
if self.llm_metrics['call_count'] > 0:
|
||||
self.log_final_metrics()
|
||||
self.reset_llm_metrics()
|
||||
self.log(f'Ending Trace for {self.event_type}')
|
||||
|
||||
self.log(f'Ending Trace for {self.event_type}', extra_fields={'trace_duration': trace_total_time})
|
||||
self._flush_log_buffer()
|
||||
|
||||
|
||||
return BusinessEventContext(self).__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
async def __aenter__(self):
|
||||
self.trace_start_time = time.time()
|
||||
self.log(f'Starting Trace for {self.event_type}')
|
||||
return await BusinessEventContext(self).__aenter__()
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
trace_total_time = time.time() - self.trace_start_time
|
||||
|
||||
# Record trace duration
|
||||
TRACE_DURATION.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).observe(trace_total_time)
|
||||
|
||||
# Decrement concurrent traces gauge
|
||||
CONCURRENT_TRACES.labels(
|
||||
tenant_id=self.tenant_id_str,
|
||||
event_type=self.event_type_str,
|
||||
specialist_id=self.specialist_id_str,
|
||||
specialist_type=self.specialist_type_str,
|
||||
specialist_type_version=self.specialist_type_version_str
|
||||
).dec()
|
||||
|
||||
self._push_to_gateway()
|
||||
|
||||
if self.llm_metrics['call_count'] > 0:
|
||||
self.log_final_metrics()
|
||||
self.reset_llm_metrics()
|
||||
|
||||
self.log(f'Ending Trace for {self.event_type}', extra_fields={'trace_duration': trace_total_time})
|
||||
self._flush_log_buffer()
|
||||
return await BusinessEventContext(self).__aexit__(exc_type, exc_val, exc_tb)
|
||||
@@ -1,9 +1,22 @@
|
||||
from werkzeug.local import LocalProxy, LocalStack
|
||||
import asyncio
|
||||
from contextvars import ContextVar
|
||||
import contextvars
|
||||
|
||||
# Keep existing stack for backward compatibility
|
||||
_business_event_stack = LocalStack()
|
||||
|
||||
# Add contextvar for async support
|
||||
_business_event_contextvar = ContextVar('business_event', default=None)
|
||||
|
||||
|
||||
def _get_current_event():
|
||||
# Try contextvar first (for async)
|
||||
event = _business_event_contextvar.get()
|
||||
if event is not None:
|
||||
return event
|
||||
|
||||
# Fall back to the stack-based approach (for sync)
|
||||
top = _business_event_stack.top
|
||||
if top is None:
|
||||
raise RuntimeError("No business event context found. Are you sure you're in a business event?")
|
||||
@@ -16,10 +29,24 @@ current_event = LocalProxy(_get_current_event)
|
||||
class BusinessEventContext:
|
||||
def __init__(self, event):
|
||||
self.event = event
|
||||
self._token = None # For storing contextvar token
|
||||
|
||||
def __enter__(self):
|
||||
_business_event_stack.push(self.event)
|
||||
self._token = _business_event_contextvar.set(self.event)
|
||||
return self.event
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
_business_event_stack.pop()
|
||||
if self._token is not None:
|
||||
_business_event_contextvar.reset(self._token)
|
||||
|
||||
async def __aenter__(self):
|
||||
_business_event_stack.push(self.event)
|
||||
self._token = _business_event_contextvar.set(self.event)
|
||||
return self.event
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
_business_event_stack.pop()
|
||||
if self._token is not None:
|
||||
_business_event_contextvar.reset(self._token)
|
||||
192
common/utils/cache/base.py
vendored
Normal file
192
common/utils/cache/base.py
vendored
Normal file
@@ -0,0 +1,192 @@
|
||||
from typing import Any, Dict, List, Optional, TypeVar, Generic, Type
|
||||
from dataclasses import dataclass
|
||||
from flask import Flask, current_app
|
||||
from dogpile.cache import CacheRegion
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
T = TypeVar('T') # Generic type parameter for cached data
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheKey:
|
||||
"""
|
||||
Represents a composite cache key made up of multiple components.
|
||||
Enables structured and consistent key generation for cache entries.
|
||||
|
||||
Attributes:
|
||||
components (Dict[str, Any]): Dictionary of key components and their values
|
||||
|
||||
Example:
|
||||
key = CacheKey({'tenant_id': 123, 'user_id': 456})
|
||||
str(key) -> "tenant_id=123:user_id=456"
|
||||
"""
|
||||
components: Dict[str, Any]
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
Converts components into a deterministic string representation.
|
||||
Components are sorted alphabetically to ensure consistent key generation.
|
||||
"""
|
||||
return ":".join(f"{k}={v}" for k, v in sorted(self.components.items()))
|
||||
|
||||
|
||||
class CacheHandler(Generic[T]):
|
||||
"""
|
||||
Base cache handler implementation providing structured caching functionality.
|
||||
Uses generics to ensure type safety of cached data.
|
||||
|
||||
Type Parameters:
|
||||
T: Type of data being cached
|
||||
|
||||
Attributes:
|
||||
region (CacheRegion): Dogpile cache region for storage
|
||||
prefix (str): Prefix for all cache keys managed by this handler
|
||||
"""
|
||||
|
||||
def __init__(self, region: CacheRegion, prefix: str):
|
||||
self.region = region
|
||||
self.prefix = prefix
|
||||
self._key_components = [] # List of required key components
|
||||
|
||||
@abstractmethod
|
||||
def _to_cache_data(self, instance: T) -> Any:
|
||||
"""
|
||||
Convert the data to a cacheable format for internal use.
|
||||
|
||||
Args:
|
||||
instance: The data to be cached.
|
||||
|
||||
Returns:
|
||||
A serializable format of the instance.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _from_cache_data(self, data: Any, **kwargs) -> T:
|
||||
"""
|
||||
Convert cached data back to usable format for internal use.
|
||||
|
||||
Args:
|
||||
data: The cached data.
|
||||
**kwargs: Additional context.
|
||||
|
||||
Returns:
|
||||
The data in its usable format.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _should_cache(self, value: T) -> bool:
|
||||
"""
|
||||
Validate if the value should be cached for internal use.
|
||||
|
||||
Args:
|
||||
value: The value to be cached.
|
||||
|
||||
Returns:
|
||||
True if the value should be cached, False otherwise.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def configure_keys(self, *components: str):
|
||||
"""
|
||||
Configure required components for cache key generation.
|
||||
|
||||
Args:
|
||||
*components: Required key component names
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
self._key_components = components
|
||||
return self
|
||||
|
||||
def generate_key(self, **identifiers) -> str:
|
||||
"""
|
||||
Generate a cache key from provided identifiers.
|
||||
|
||||
Args:
|
||||
**identifiers: Key-value pairs for key components
|
||||
|
||||
Returns:
|
||||
Formatted cache key string
|
||||
|
||||
Raises:
|
||||
ValueError: If required components are missing
|
||||
"""
|
||||
missing = set(self._key_components) - set(identifiers.keys())
|
||||
if missing:
|
||||
raise ValueError(f"Missing key components: {missing}")
|
||||
|
||||
region_name = getattr(self.region, 'name', 'default_region')
|
||||
|
||||
key = CacheKey({k: identifiers[k] for k in self._key_components})
|
||||
return f"{region_name}_{self.prefix}:{str(key)}"
|
||||
|
||||
def get(self, creator_func, **identifiers) -> T:
|
||||
"""
|
||||
Get or create a cached value.
|
||||
|
||||
Args:
|
||||
creator_func: Function to create value if not cached
|
||||
**identifiers: Key components for cache key
|
||||
|
||||
Returns:
|
||||
Cached or newly created value
|
||||
"""
|
||||
cache_key = self.generate_key(**identifiers)
|
||||
|
||||
def creator():
|
||||
instance = creator_func(**identifiers)
|
||||
serialized_instance = self._to_cache_data(instance)
|
||||
return serialized_instance
|
||||
|
||||
cached_data = self.region.get_or_create(
|
||||
cache_key,
|
||||
creator,
|
||||
should_cache_fn=self._should_cache
|
||||
)
|
||||
|
||||
return self._from_cache_data(cached_data, **identifiers)
|
||||
|
||||
def invalidate(self, **identifiers):
|
||||
"""
|
||||
Invalidate a specific cache entry.
|
||||
|
||||
Args:
|
||||
**identifiers: Key components for the cache entry
|
||||
"""
|
||||
cache_key = self.generate_key(**identifiers)
|
||||
self.region.delete(cache_key)
|
||||
|
||||
def invalidate_by_model(self, model: str, **identifiers):
|
||||
"""
|
||||
Invalidate cache entry based on model changes.
|
||||
|
||||
Args:
|
||||
model: Changed model name
|
||||
**identifiers: Model instance identifiers
|
||||
"""
|
||||
try:
|
||||
self.invalidate(**identifiers)
|
||||
except ValueError:
|
||||
pass # Skip if cache key can't be generated from provided identifiers
|
||||
|
||||
def invalidate_region(self):
|
||||
"""
|
||||
Invalidate all cache entries within this region.
|
||||
|
||||
Deletes all keys that start with the region prefix.
|
||||
"""
|
||||
# Construct the pattern for all keys in this region
|
||||
pattern = f"{self.region}_{self.prefix}:*"
|
||||
|
||||
# Assuming Redis backend with dogpile, use `delete_multi` or direct Redis access
|
||||
if hasattr(self.region.backend, 'client'):
|
||||
redis_client = self.region.backend.client
|
||||
keys_to_delete = redis_client.keys(pattern)
|
||||
if keys_to_delete:
|
||||
redis_client.delete(*keys_to_delete)
|
||||
else:
|
||||
# Fallback for other backends
|
||||
raise NotImplementedError("Region invalidation is only supported for Redis backend.")
|
||||
515
common/utils/cache/config_cache.py
vendored
Normal file
515
common/utils/cache/config_cache.py
vendored
Normal file
@@ -0,0 +1,515 @@
|
||||
from typing import Dict, Any, Optional
|
||||
from pathlib import Path
|
||||
import yaml
|
||||
from packaging import version
|
||||
import os
|
||||
from flask import current_app
|
||||
|
||||
from common.utils.cache.base import CacheHandler, CacheKey
|
||||
from config.type_defs import agent_types, task_types, tool_types, specialist_types, retriever_types, prompt_types, \
|
||||
catalog_types, partner_service_types, processor_types
|
||||
|
||||
|
||||
def is_major_minor(version: str) -> bool:
|
||||
parts = version.strip('.').split('.')
|
||||
return len(parts) == 2 and all(part.isdigit() for part in parts)
|
||||
|
||||
|
||||
class BaseConfigCacheHandler(CacheHandler[Dict[str, Any]]):
|
||||
"""Base handler for configuration caching"""
|
||||
|
||||
def __init__(self, region, config_type: str):
|
||||
"""
|
||||
Args:
|
||||
region: Cache region
|
||||
config_type: Type of configuration (agents, tasks, etc.)
|
||||
"""
|
||||
super().__init__(region, f'config_{config_type}')
|
||||
self.config_type = config_type
|
||||
self._types_module = None # Set by subclasses
|
||||
self._config_dir = None # Set by subclasses
|
||||
self.version_tree_cache = None
|
||||
self.configure_keys('type_name', 'version')
|
||||
|
||||
def _to_cache_data(self, instance: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Convert the data to a cacheable format"""
|
||||
# For configuration data, we can just return the dictionary as is
|
||||
# since it's already in a serializable format
|
||||
return instance
|
||||
|
||||
def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
|
||||
"""Convert cached data back to usable format"""
|
||||
# Similarly, we can return the data directly since it's already
|
||||
# in the format we need
|
||||
return data
|
||||
|
||||
def _should_cache(self, value: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate if the value should be cached
|
||||
|
||||
Args:
|
||||
value: The value to be cached
|
||||
|
||||
Returns:
|
||||
bool: True if the value should be cached
|
||||
"""
|
||||
return isinstance(value, dict) # Cache all dictionaries
|
||||
|
||||
def set_version_tree_cache(self, cache):
|
||||
"""Set the version tree cache dependency."""
|
||||
self.version_tree_cache = cache
|
||||
|
||||
def _load_specific_config(self, type_name: str, version_str: str = 'latest') -> Dict[str, Any]:
|
||||
"""
|
||||
Load a specific configuration version
|
||||
Automatically handles global vs partner-specific configs
|
||||
"""
|
||||
version_tree = self.version_tree_cache.get_versions(type_name)
|
||||
versions = version_tree['versions']
|
||||
|
||||
if version_str == 'latest':
|
||||
version_str = version_tree['latest_version']
|
||||
|
||||
if version_str not in versions:
|
||||
raise ValueError(f"Version {version_str} not found for {type_name}")
|
||||
|
||||
version_info = versions[version_str]
|
||||
file_path = version_info['file_path']
|
||||
partner = version_info.get('partner')
|
||||
|
||||
try:
|
||||
with open(file_path) as f:
|
||||
config = yaml.safe_load(f)
|
||||
# Add partner information to the config
|
||||
if partner:
|
||||
config['partner'] = partner
|
||||
return config
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error loading config from {file_path}: {e}")
|
||||
|
||||
def get_config(self, type_name: str, version: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""
|
||||
Get configuration for a specific type and version
|
||||
If version not specified, returns latest
|
||||
|
||||
Args:
|
||||
type_name: Configuration type name
|
||||
version: Optional specific version to retrieve
|
||||
|
||||
Returns:
|
||||
Configuration data
|
||||
"""
|
||||
if version is None:
|
||||
version_str = self.version_tree_cache.get_latest_version(type_name)
|
||||
elif is_major_minor(version):
|
||||
version_str = self.version_tree_cache.get_latest_patch_version(type_name, version)
|
||||
else:
|
||||
version_str = version
|
||||
|
||||
result = self.get(
|
||||
lambda type_name, version: self._load_specific_config(type_name, version),
|
||||
type_name=type_name,
|
||||
version=version_str
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class BaseConfigVersionTreeCacheHandler(CacheHandler[Dict[str, Any]]):
|
||||
"""Base handler for configuration version tree caching"""
|
||||
|
||||
def __init__(self, region, config_type: str):
|
||||
"""
|
||||
Args:
|
||||
region: Cache region
|
||||
config_type: Type of configuration (agents, tasks, etc.)
|
||||
"""
|
||||
super().__init__(region, f'config_{config_type}_version_tree')
|
||||
self.config_type = config_type
|
||||
self._types_module = None # Set by subclasses
|
||||
self._config_dir = None # Set by subclasses
|
||||
self.configure_keys('type_name')
|
||||
|
||||
def _load_version_tree(self, type_name: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Load version tree for a specific type without loading full configurations
|
||||
Checks both global and partner-specific directories
|
||||
"""
|
||||
# First check the global path
|
||||
global_path = Path(self._config_dir) / "globals" / type_name
|
||||
|
||||
# If global path doesn't exist, check if the type exists directly in the root
|
||||
# (for backward compatibility)
|
||||
if not global_path.exists():
|
||||
global_path = Path(self._config_dir) / type_name
|
||||
|
||||
if not global_path.exists():
|
||||
# Check if it exists in any partner subdirectories
|
||||
partner_dirs = [d for d in Path(self._config_dir).iterdir()
|
||||
if d.is_dir() and d.name != "globals"]
|
||||
|
||||
for partner_dir in partner_dirs:
|
||||
partner_type_path = partner_dir / type_name
|
||||
if partner_type_path.exists():
|
||||
# Found in partner directory
|
||||
return self._load_versions_from_path(partner_type_path)
|
||||
|
||||
# If we get here, the type wasn't found anywhere
|
||||
raise ValueError(f"No configuration found for type {type_name}")
|
||||
|
||||
return self._load_versions_from_path(global_path)
|
||||
|
||||
def _load_versions_from_path(self, path: Path) -> Dict[str, Any]:
|
||||
"""Load all versions from a specific path"""
|
||||
version_files = list(path.glob('*.yaml'))
|
||||
if not version_files:
|
||||
raise ValueError(f"No versions found in {path}")
|
||||
|
||||
versions = {}
|
||||
latest_version = None
|
||||
latest_version_obj = None
|
||||
|
||||
for file_path in version_files:
|
||||
ver = file_path.stem # Get version from filename
|
||||
try:
|
||||
ver_obj = version.parse(ver)
|
||||
# Only load minimal metadata for version tree
|
||||
with open(file_path) as f:
|
||||
yaml_data = yaml.safe_load(f)
|
||||
metadata = yaml_data.get('metadata', {})
|
||||
# Add partner information if available
|
||||
partner = None
|
||||
if "globals" not in str(file_path):
|
||||
# Extract partner name from path
|
||||
# Path format: config_dir/partner_name/type_name/version.yaml
|
||||
partner = file_path.parent.parent.name
|
||||
|
||||
versions[ver] = {
|
||||
'metadata': metadata,
|
||||
'file_path': str(file_path),
|
||||
'partner': partner
|
||||
}
|
||||
|
||||
# Track latest version
|
||||
if latest_version_obj is None or ver_obj > latest_version_obj:
|
||||
latest_version = ver
|
||||
latest_version_obj = ver_obj
|
||||
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error loading version {ver}: {e}")
|
||||
continue
|
||||
|
||||
return {
|
||||
'versions': versions,
|
||||
'latest_version': latest_version
|
||||
}
|
||||
|
||||
def _to_cache_data(self, instance: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Convert the data to a cacheable format"""
|
||||
# For configuration data, we can just return the dictionary as is
|
||||
# since it's already in a serializable format
|
||||
return instance
|
||||
|
||||
def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
|
||||
"""Convert cached data back to usable format"""
|
||||
# Similarly, we can return the data directly since it's already
|
||||
# in the format we need
|
||||
return data
|
||||
|
||||
def _should_cache(self, value: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate if the value should be cached
|
||||
|
||||
Args:
|
||||
value: The value to be cached
|
||||
|
||||
Returns:
|
||||
bool: True if the value should be cached
|
||||
"""
|
||||
return isinstance(value, dict) # Cache all dictionaries
|
||||
|
||||
def get_versions(self, type_name: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Get version tree for a type
|
||||
|
||||
Args:
|
||||
type_name: Type to get versions for
|
||||
|
||||
Returns:
|
||||
Dict with version information
|
||||
"""
|
||||
return self.get(
|
||||
lambda type_name: self._load_version_tree(type_name),
|
||||
type_name=type_name
|
||||
)
|
||||
|
||||
def get_latest_version(self, type_name: str) -> str:
|
||||
"""
|
||||
Get the latest version for a given type name.
|
||||
|
||||
Args:
|
||||
type_name: Name of the configuration type
|
||||
|
||||
Returns:
|
||||
Latest version string
|
||||
|
||||
Raises:
|
||||
ValueError: If type not found or no versions available
|
||||
"""
|
||||
version_tree = self.get_versions(type_name)
|
||||
if not version_tree or 'latest_version' not in version_tree:
|
||||
raise ValueError(f"No versions found for {type_name}")
|
||||
|
||||
return version_tree['latest_version']
|
||||
|
||||
def get_latest_patch_version(self, type_name: str, major_minor: str) -> str:
|
||||
"""
|
||||
Get the latest patch version for a given major.minor version.
|
||||
|
||||
Args:
|
||||
type_name: Name of the configuration type
|
||||
major_minor: Major.minor version (e.g. "1.0")
|
||||
|
||||
Returns:
|
||||
Latest patch version string (e.g. "1.0.3")
|
||||
|
||||
Raises:
|
||||
ValueError: If type not found or no matching versions
|
||||
"""
|
||||
version_tree = self.get_versions(type_name)
|
||||
if not version_tree or 'versions' not in version_tree:
|
||||
raise ValueError(f"No versions found for {type_name}")
|
||||
|
||||
# Filter versions that match the major.minor prefix
|
||||
matching_versions = [
|
||||
ver for ver in version_tree['versions'].keys()
|
||||
if ver.startswith(major_minor + '.')
|
||||
]
|
||||
|
||||
if not matching_versions:
|
||||
raise ValueError(f"No versions found for {type_name} with prefix {major_minor}")
|
||||
|
||||
# Return highest matching version
|
||||
latest_patch = max(matching_versions, key=version.parse)
|
||||
return latest_patch
|
||||
|
||||
|
||||
class BaseConfigTypesCacheHandler(CacheHandler[Dict[str, Any]]):
|
||||
"""Base handler for configuration types caching"""
|
||||
|
||||
def __init__(self, region, config_type: str):
|
||||
"""
|
||||
Args:
|
||||
region: Cache region
|
||||
config_type: Type of configuration (agents, tasks, etc.)
|
||||
"""
|
||||
super().__init__(region, f'config_{config_type}_types')
|
||||
self.config_type = config_type
|
||||
self._types_module = None # Set by subclasses
|
||||
self._config_dir = None # Set by subclasses
|
||||
self.configure_keys()
|
||||
|
||||
def _to_cache_data(self, instance: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Convert the data to a cacheable format"""
|
||||
# For configuration data, we can just return the dictionary as is
|
||||
# since it's already in a serializable format
|
||||
return instance
|
||||
|
||||
def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
|
||||
"""Convert cached data back to usable format"""
|
||||
# Similarly, we can return the data directly since it's already
|
||||
# in the format we need
|
||||
return data
|
||||
|
||||
def _should_cache(self, value: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate if the value should be cached
|
||||
|
||||
Args:
|
||||
value: The value to be cached
|
||||
|
||||
Returns:
|
||||
bool: True if the value should be cached
|
||||
"""
|
||||
return isinstance(value, dict) # Cache all dictionaries
|
||||
|
||||
def _load_type_definitions(self) -> Dict[str, Dict[str, str]]:
|
||||
"""Load type definitions from the corresponding type_defs module"""
|
||||
if not self._types_module:
|
||||
raise ValueError("_types_module must be set by subclass")
|
||||
|
||||
type_definitions = {
|
||||
type_id: {
|
||||
'name': info['name'],
|
||||
'description': info['description'],
|
||||
'partner': info.get('partner') # Include partner info if available
|
||||
}
|
||||
for type_id, info in self._types_module.items()
|
||||
}
|
||||
|
||||
return type_definitions
|
||||
|
||||
def get_types(self) -> Dict[str, Dict[str, str]]:
|
||||
"""Get dictionary of available types with name and description"""
|
||||
result = self.get(
|
||||
lambda type_name: self._load_type_definitions(),
|
||||
type_name=f'{self.config_type}_types',
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def create_config_cache_handlers(config_type: str, config_dir: str, types_module: dict) -> tuple:
|
||||
"""
|
||||
Factory function to dynamically create the 3 cache handler classes for a given configuration type.
|
||||
The following cache names are created:
|
||||
- <config_type>_config_cache
|
||||
- <config_type>_version_tree_cache
|
||||
- <config_type>_types_cache
|
||||
|
||||
|
||||
Args:
|
||||
config_type: The configuration type (e.g., 'agents', 'tasks').
|
||||
config_dir: The directory where configuration files are stored.
|
||||
types_module: The types module defining the available types for this config.
|
||||
|
||||
Returns:
|
||||
A tuple of dynamically created classes for config, version tree, and types handlers.
|
||||
"""
|
||||
|
||||
class ConfigCacheHandler(BaseConfigCacheHandler):
|
||||
handler_name = f"{config_type}_config_cache"
|
||||
|
||||
def __init__(self, region):
|
||||
super().__init__(region, config_type)
|
||||
self._types_module = types_module
|
||||
self._config_dir = config_dir
|
||||
|
||||
class VersionTreeCacheHandler(BaseConfigVersionTreeCacheHandler):
|
||||
handler_name = f"{config_type}_version_tree_cache"
|
||||
|
||||
def __init__(self, region):
|
||||
super().__init__(region, config_type)
|
||||
self._types_module = types_module
|
||||
self._config_dir = config_dir
|
||||
|
||||
class TypesCacheHandler(BaseConfigTypesCacheHandler):
|
||||
handler_name = f"{config_type}_types_cache"
|
||||
|
||||
def __init__(self, region):
|
||||
super().__init__(region, config_type)
|
||||
self._types_module = types_module
|
||||
self._config_dir = config_dir
|
||||
|
||||
return ConfigCacheHandler, VersionTreeCacheHandler, TypesCacheHandler
|
||||
|
||||
|
||||
AgentConfigCacheHandler, AgentConfigVersionTreeCacheHandler, AgentConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='agents',
|
||||
config_dir='config/agents',
|
||||
types_module=agent_types.AGENT_TYPES
|
||||
))
|
||||
|
||||
|
||||
TaskConfigCacheHandler, TaskConfigVersionTreeCacheHandler, TaskConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='tasks',
|
||||
config_dir='config/tasks',
|
||||
types_module=task_types.TASK_TYPES
|
||||
))
|
||||
|
||||
|
||||
ToolConfigCacheHandler, ToolConfigVersionTreeCacheHandler, ToolConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='tools',
|
||||
config_dir='config/tools',
|
||||
types_module=tool_types.TOOL_TYPES
|
||||
))
|
||||
|
||||
|
||||
SpecialistConfigCacheHandler, SpecialistConfigVersionTreeCacheHandler, SpecialistConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='specialists',
|
||||
config_dir='config/specialists',
|
||||
types_module=specialist_types.SPECIALIST_TYPES
|
||||
))
|
||||
|
||||
|
||||
RetrieverConfigCacheHandler, RetrieverConfigVersionTreeCacheHandler, RetrieverConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='retrievers',
|
||||
config_dir='config/retrievers',
|
||||
types_module=retriever_types.RETRIEVER_TYPES
|
||||
))
|
||||
|
||||
|
||||
PromptConfigCacheHandler, PromptConfigVersionTreeCacheHandler, PromptConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='prompts',
|
||||
config_dir='config/prompts',
|
||||
types_module=prompt_types.PROMPT_TYPES
|
||||
))
|
||||
|
||||
CatalogConfigCacheHandler, CatalogConfigVersionTreeCacheHandler, CatalogConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='catalogs',
|
||||
config_dir='config/catalogs',
|
||||
types_module=catalog_types.CATALOG_TYPES
|
||||
))
|
||||
|
||||
ProcessorConfigCacheHandler, ProcessorConfigVersionTreeCacheHandler, ProcessorConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='processors',
|
||||
config_dir='config/processors',
|
||||
types_module=processor_types.PROCESSOR_TYPES
|
||||
))
|
||||
|
||||
# Add to common/utils/cache/config_cache.py
|
||||
PartnerServiceConfigCacheHandler, PartnerServiceConfigVersionTreeCacheHandler, PartnerServiceConfigTypesCacheHandler = (
|
||||
create_config_cache_handlers(
|
||||
config_type='partner_services',
|
||||
config_dir='config/partner_services',
|
||||
types_module=partner_service_types.PARTNER_SERVICE_TYPES
|
||||
))
|
||||
|
||||
|
||||
def register_config_cache_handlers(cache_manager) -> None:
|
||||
cache_manager.register_handler(AgentConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(AgentConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(AgentConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(TaskConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(TaskConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(TaskConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(ToolConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(ToolConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(ToolConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(SpecialistConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(SpecialistConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(SpecialistConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(RetrieverConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(RetrieverConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(RetrieverConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(PromptConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(PromptConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(PromptConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(CatalogConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(CatalogConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(CatalogConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(ProcessorConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(ProcessorConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(ProcessorConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(AgentConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(AgentConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(AgentConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(PartnerServiceConfigCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(PartnerServiceConfigTypesCacheHandler, 'eveai_config')
|
||||
cache_manager.register_handler(PartnerServiceConfigVersionTreeCacheHandler, 'eveai_config')
|
||||
|
||||
cache_manager.agents_config_cache.set_version_tree_cache(cache_manager.agents_version_tree_cache)
|
||||
cache_manager.tasks_config_cache.set_version_tree_cache(cache_manager.tasks_version_tree_cache)
|
||||
cache_manager.tools_config_cache.set_version_tree_cache(cache_manager.tools_version_tree_cache)
|
||||
cache_manager.specialists_config_cache.set_version_tree_cache(cache_manager.specialists_version_tree_cache)
|
||||
cache_manager.retrievers_config_cache.set_version_tree_cache(cache_manager.retrievers_version_tree_cache)
|
||||
cache_manager.prompts_config_cache.set_version_tree_cache(cache_manager.prompts_version_tree_cache)
|
||||
cache_manager.catalogs_config_cache.set_version_tree_cache(cache_manager.catalogs_version_tree_cache)
|
||||
cache_manager.processors_config_cache.set_version_tree_cache(cache_manager.processors_version_tree_cache)
|
||||
cache_manager.partner_services_config_cache.set_version_tree_cache(cache_manager.partner_services_version_tree_cache)
|
||||
218
common/utils/cache/crewai_config_processor.py
vendored
Normal file
218
common/utils/cache/crewai_config_processor.py
vendored
Normal file
@@ -0,0 +1,218 @@
|
||||
from typing import Dict, Any, Type, TypeVar, List
|
||||
from abc import ABC, abstractmethod
|
||||
from flask import current_app
|
||||
|
||||
from common.extensions import cache_manager, db
|
||||
from common.models.interaction import EveAIAgent, EveAITask, EveAITool, Specialist
|
||||
from common.utils.cache.crewai_configuration import (
|
||||
ProcessedAgentConfig, ProcessedTaskConfig, ProcessedToolConfig,
|
||||
SpecialistProcessedConfig
|
||||
)
|
||||
|
||||
T = TypeVar('T') # For generic model types
|
||||
|
||||
|
||||
class BaseCrewAIConfigProcessor:
|
||||
"""Base processor for specialist configurations"""
|
||||
|
||||
# Standard mapping between model fields and template placeholders
|
||||
AGENT_FIELD_MAPPING = {
|
||||
'role': 'custom_role',
|
||||
'goal': 'custom_goal',
|
||||
'backstory': 'custom_backstory'
|
||||
}
|
||||
|
||||
TASK_FIELD_MAPPING = {
|
||||
'task_description': 'custom_description',
|
||||
'expected_output': 'custom_expected_output'
|
||||
}
|
||||
|
||||
def __init__(self, tenant_id: int, specialist_id: int):
|
||||
self.tenant_id = tenant_id
|
||||
self.specialist_id = specialist_id
|
||||
self.specialist = self._get_specialist()
|
||||
self.verbose = self._get_verbose_setting()
|
||||
|
||||
def _get_specialist(self) -> Specialist:
|
||||
"""Get specialist and verify existence"""
|
||||
specialist = Specialist.query.get(self.specialist_id)
|
||||
if not specialist:
|
||||
raise ValueError(f"Specialist {self.specialist_id} not found")
|
||||
return specialist
|
||||
|
||||
def _get_verbose_setting(self) -> bool:
|
||||
"""Get verbose setting from specialist"""
|
||||
return bool(self.specialist.tuning)
|
||||
|
||||
def _get_db_items(self, model_class: Type[T], type_list: List[str]) -> Dict[str, T]:
|
||||
"""Get database items of specified type"""
|
||||
items = (model_class.query
|
||||
.filter_by(specialist_id=self.specialist_id)
|
||||
.filter(model_class.type.in_(type_list))
|
||||
.all())
|
||||
return {item.type: item for item in items}
|
||||
|
||||
def _apply_replacements(self, text: str, replacements: Dict[str, str]) -> str:
|
||||
"""Apply text replacements to a string"""
|
||||
result = text
|
||||
for key, value in replacements.items():
|
||||
if value is not None: # Only replace if value exists
|
||||
placeholder = "{" + key + "}"
|
||||
result = result.replace(placeholder, str(value))
|
||||
return result
|
||||
|
||||
def _process_agent_configs(self, specialist_config: Dict[str, Any]) -> Dict[str, ProcessedAgentConfig]:
|
||||
"""Process all agent configurations"""
|
||||
agent_configs = {}
|
||||
|
||||
if 'agents' not in specialist_config:
|
||||
return agent_configs
|
||||
|
||||
# Get all DB agents at once
|
||||
agent_types = [agent_def['type'] for agent_def in specialist_config['agents']]
|
||||
db_agents = self._get_db_items(EveAIAgent, agent_types)
|
||||
|
||||
for agent_def in specialist_config['agents']:
|
||||
agent_type = agent_def['type']
|
||||
agent_type_lower = agent_type.lower()
|
||||
db_agent = db_agents.get(agent_type)
|
||||
|
||||
# Get full configuration
|
||||
config = cache_manager.agents_config_cache.get_config(
|
||||
agent_type,
|
||||
agent_def.get('version', '1.0')
|
||||
)
|
||||
|
||||
# Start with YAML values
|
||||
role = config['role']
|
||||
goal = config['goal']
|
||||
backstory = config['backstory']
|
||||
|
||||
# Apply DB values if they exist
|
||||
if db_agent:
|
||||
for model_field, placeholder in self.AGENT_FIELD_MAPPING.items():
|
||||
value = getattr(db_agent, model_field)
|
||||
if value:
|
||||
placeholder_text = "{" + placeholder + "}"
|
||||
role = role.replace(placeholder_text, value)
|
||||
goal = goal.replace(placeholder_text, value)
|
||||
backstory = backstory.replace(placeholder_text, value)
|
||||
|
||||
agent_configs[agent_type_lower] = ProcessedAgentConfig(
|
||||
role=role,
|
||||
goal=goal,
|
||||
backstory=backstory,
|
||||
name=agent_def.get('name') or config.get('name', agent_type_lower),
|
||||
type=agent_type,
|
||||
description=agent_def.get('description') or config.get('description'),
|
||||
verbose=self.verbose
|
||||
)
|
||||
|
||||
return agent_configs
|
||||
|
||||
def _process_task_configs(self, specialist_config: Dict[str, Any]) -> Dict[str, ProcessedTaskConfig]:
|
||||
"""Process all task configurations"""
|
||||
task_configs = {}
|
||||
|
||||
if 'tasks' not in specialist_config:
|
||||
return task_configs
|
||||
|
||||
# Get all DB tasks at once
|
||||
task_types = [task_def['type'] for task_def in specialist_config['tasks']]
|
||||
db_tasks = self._get_db_items(EveAITask, task_types)
|
||||
|
||||
for task_def in specialist_config['tasks']:
|
||||
task_type = task_def['type']
|
||||
task_type_lower = task_type.lower()
|
||||
db_task = db_tasks.get(task_type)
|
||||
|
||||
# Get full configuration
|
||||
config = cache_manager.tasks_config_cache.get_config(
|
||||
task_type,
|
||||
task_def.get('version', '1.0')
|
||||
)
|
||||
|
||||
# Start with YAML values
|
||||
task_description = config['task_description']
|
||||
expected_output = config['expected_output']
|
||||
|
||||
# Apply DB values if they exist
|
||||
if db_task:
|
||||
for model_field, placeholder in self.TASK_FIELD_MAPPING.items():
|
||||
value = getattr(db_task, model_field)
|
||||
if value:
|
||||
placeholder_text = "{" + placeholder + "}"
|
||||
task_description = task_description.replace(placeholder_text, value)
|
||||
expected_output = expected_output.replace(placeholder_text, value)
|
||||
|
||||
task_configs[task_type_lower] = ProcessedTaskConfig(
|
||||
task_description=task_description,
|
||||
expected_output=expected_output,
|
||||
name=task_def.get('name') or config.get('name', task_type_lower),
|
||||
type=task_type,
|
||||
description=task_def.get('description') or config.get('description'),
|
||||
verbose=self.verbose
|
||||
)
|
||||
|
||||
return task_configs
|
||||
|
||||
def _process_tool_configs(self, specialist_config: Dict[str, Any]) -> Dict[str, ProcessedToolConfig]:
|
||||
"""Process all tool configurations"""
|
||||
tool_configs = {}
|
||||
|
||||
if 'tools' not in specialist_config:
|
||||
return tool_configs
|
||||
|
||||
# Get all DB tools at once
|
||||
tool_types = [tool_def['type'] for tool_def in specialist_config['tools']]
|
||||
db_tools = self._get_db_items(EveAITool, tool_types)
|
||||
|
||||
for tool_def in specialist_config['tools']:
|
||||
tool_type = tool_def['type']
|
||||
tool_type_lower = tool_type.lower()
|
||||
db_tool = db_tools.get(tool_type)
|
||||
|
||||
# Get full configuration
|
||||
config = cache_manager.tools_config_cache.get_config(
|
||||
tool_type,
|
||||
tool_def.get('version', '1.0')
|
||||
)
|
||||
|
||||
# Combine configuration
|
||||
tool_config = config.get('configuration', {})
|
||||
if db_tool and db_tool.configuration:
|
||||
tool_config.update(db_tool.configuration)
|
||||
|
||||
tool_configs[tool_type_lower] = ProcessedToolConfig(
|
||||
name=tool_def.get('name') or config.get('name', tool_type_lower),
|
||||
type=tool_type,
|
||||
description=tool_def.get('description') or config.get('description'),
|
||||
configuration=tool_config,
|
||||
verbose=self.verbose
|
||||
)
|
||||
|
||||
return tool_configs
|
||||
|
||||
def process_config(self) -> SpecialistProcessedConfig:
|
||||
"""Process complete specialist configuration"""
|
||||
try:
|
||||
# Get full specialist configuration
|
||||
specialist_config = cache_manager.specialists_config_cache.get_config(
|
||||
self.specialist.type,
|
||||
self.specialist.type_version
|
||||
)
|
||||
|
||||
if not specialist_config:
|
||||
raise ValueError(f"No configuration found for {self.specialist.type}")
|
||||
|
||||
# Process all configurations
|
||||
processed_config = SpecialistProcessedConfig(
|
||||
agents=self._process_agent_configs(specialist_config),
|
||||
tasks=self._process_task_configs(specialist_config),
|
||||
tools=self._process_tool_configs(specialist_config)
|
||||
)
|
||||
return processed_config
|
||||
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error processing specialist configuration: {e}")
|
||||
raise
|
||||
126
common/utils/cache/crewai_configuration.py
vendored
Normal file
126
common/utils/cache/crewai_configuration.py
vendored
Normal file
@@ -0,0 +1,126 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Any, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessedAgentConfig:
|
||||
"""Processed and ready-to-use agent configuration"""
|
||||
role: str
|
||||
goal: str
|
||||
backstory: str
|
||||
name: str
|
||||
type: str
|
||||
description: Optional[str] = None
|
||||
verbose: bool = False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for serialization"""
|
||||
return {
|
||||
'role': self.role,
|
||||
'goal': self.goal,
|
||||
'backstory': self.backstory,
|
||||
'name': self.name,
|
||||
'type': self.type,
|
||||
'description': self.description,
|
||||
'verbose': self.verbose
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'ProcessedAgentConfig':
|
||||
"""Create from dictionary"""
|
||||
return cls(**data)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessedTaskConfig:
|
||||
"""Processed and ready-to-use task configuration"""
|
||||
task_description: str
|
||||
expected_output: str
|
||||
name: str
|
||||
type: str
|
||||
description: Optional[str] = None
|
||||
verbose: bool = False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for serialization"""
|
||||
return {
|
||||
'task_description': self.task_description,
|
||||
'expected_output': self.expected_output,
|
||||
'name': self.name,
|
||||
'type': self.type,
|
||||
'description': self.description,
|
||||
'verbose': self.verbose
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'ProcessedTaskConfig':
|
||||
"""Create from dictionary"""
|
||||
return cls(**data)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessedToolConfig:
|
||||
"""Processed and ready-to-use tool configuration"""
|
||||
name: str
|
||||
type: str
|
||||
description: Optional[str] = None
|
||||
configuration: Optional[Dict[str, Any]] = None
|
||||
verbose: bool = False
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary for serialization"""
|
||||
return {
|
||||
'name': self.name,
|
||||
'type': self.type,
|
||||
'description': self.description,
|
||||
'configuration': self.configuration,
|
||||
'verbose': self.verbose
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'ProcessedToolConfig':
|
||||
"""Create from dictionary"""
|
||||
return cls(**data)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SpecialistProcessedConfig:
|
||||
"""Complete processed configuration for a specialist"""
|
||||
agents: Dict[str, ProcessedAgentConfig]
|
||||
tasks: Dict[str, ProcessedTaskConfig]
|
||||
tools: Dict[str, ProcessedToolConfig]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert entire configuration to dictionary"""
|
||||
return {
|
||||
'agents': {
|
||||
agent_type: config.to_dict()
|
||||
for agent_type, config in self.agents.items()
|
||||
},
|
||||
'tasks': {
|
||||
task_type: config.to_dict()
|
||||
for task_type, config in self.tasks.items()
|
||||
},
|
||||
'tools': {
|
||||
tool_type: config.to_dict()
|
||||
for tool_type, config in self.tools.items()
|
||||
}
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> 'SpecialistProcessedConfig':
|
||||
"""Create from dictionary"""
|
||||
return cls(
|
||||
agents={
|
||||
agent_type: ProcessedAgentConfig.from_dict(config)
|
||||
for agent_type, config in data['agents'].items()
|
||||
},
|
||||
tasks={
|
||||
task_type: ProcessedTaskConfig.from_dict(config)
|
||||
for task_type, config in data['tasks'].items()
|
||||
},
|
||||
tools={
|
||||
tool_type: ProcessedToolConfig.from_dict(config)
|
||||
for tool_type, config in data['tools'].items()
|
||||
}
|
||||
)
|
||||
75
common/utils/cache/crewai_processed_config_cache.py
vendored
Normal file
75
common/utils/cache/crewai_processed_config_cache.py
vendored
Normal file
@@ -0,0 +1,75 @@
|
||||
from typing import Dict, Any, Type
|
||||
from flask import current_app
|
||||
|
||||
from common.utils.cache.base import CacheHandler
|
||||
from common.utils.cache.crewai_configuration import SpecialistProcessedConfig
|
||||
from common.utils.cache.crewai_config_processor import BaseCrewAIConfigProcessor
|
||||
|
||||
|
||||
class CrewAIProcessedConfigCacheHandler(CacheHandler[SpecialistProcessedConfig]):
|
||||
"""Handles caching of processed specialist configurations"""
|
||||
handler_name = 'crewai_processed_config_cache'
|
||||
|
||||
def __init__(self, region):
|
||||
super().__init__(region, 'crewai_processed_config')
|
||||
self.configure_keys('tenant_id', 'specialist_id')
|
||||
|
||||
def _to_cache_data(self, instance: SpecialistProcessedConfig) -> Dict[str, Any]:
|
||||
"""Convert SpecialistProcessedConfig to cache data"""
|
||||
return instance.to_dict()
|
||||
|
||||
def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> SpecialistProcessedConfig:
|
||||
"""Create SpecialistProcessedConfig from cache data"""
|
||||
return SpecialistProcessedConfig.from_dict(data)
|
||||
|
||||
def _should_cache(self, value: Dict[str, Any]) -> bool:
|
||||
"""Validate cache data"""
|
||||
required_keys = {'agents', 'tasks', 'tools'}
|
||||
if not all(key in value for key in required_keys):
|
||||
current_app.logger.warning(f'CrewAI Processed Config Cache missing required keys: {required_keys}')
|
||||
return False
|
||||
return bool(value['agents'] or value['tasks'])
|
||||
|
||||
def get_specialist_config(self, tenant_id: int, specialist_id: int) -> SpecialistProcessedConfig:
|
||||
"""
|
||||
Get or create processed configuration for a specialist
|
||||
|
||||
Args:
|
||||
tenant_id: Tenant ID
|
||||
specialist_id: Specialist ID
|
||||
|
||||
Returns:
|
||||
Processed specialist configuration
|
||||
|
||||
Raises:
|
||||
ValueError: If specialist not found or processor not configured
|
||||
"""
|
||||
|
||||
def creator_func(tenant_id: int, specialist_id: int) -> SpecialistProcessedConfig:
|
||||
# Create processor instance and process config
|
||||
processor = BaseCrewAIConfigProcessor(tenant_id, specialist_id)
|
||||
return processor.process_config()
|
||||
|
||||
return self.get(
|
||||
creator_func,
|
||||
tenant_id=tenant_id,
|
||||
specialist_id=specialist_id
|
||||
)
|
||||
|
||||
def invalidate_tenant_specialist(self, tenant_id: int, specialist_id: int):
|
||||
"""Invalidate cache for a specific tenant's specialist"""
|
||||
self.invalidate(
|
||||
tenant_id=tenant_id,
|
||||
specialist_id=specialist_id
|
||||
)
|
||||
current_app.logger.info(
|
||||
f"Invalidated cache for tenant {tenant_id} specialist {specialist_id}"
|
||||
)
|
||||
|
||||
|
||||
def register_specialist_cache_handlers(cache_manager) -> None:
|
||||
"""Register specialist cache handlers with cache manager"""
|
||||
cache_manager.register_handler(
|
||||
CrewAIProcessedConfigCacheHandler,
|
||||
'eveai_chat_workers'
|
||||
)
|
||||
51
common/utils/cache/eveai_cache_manager.py
vendored
Normal file
51
common/utils/cache/eveai_cache_manager.py
vendored
Normal file
@@ -0,0 +1,51 @@
|
||||
from typing import Type
|
||||
|
||||
from flask import Flask
|
||||
|
||||
from common.utils.cache.base import CacheHandler
|
||||
from common.utils.cache.regions import create_cache_regions
|
||||
from common.utils.cache.config_cache import AgentConfigCacheHandler
|
||||
|
||||
|
||||
class EveAICacheManager:
|
||||
"""Cache manager with registration capabilities"""
|
||||
|
||||
def __init__(self):
|
||||
self._regions = {}
|
||||
self._handlers = {}
|
||||
self._handler_instances = {}
|
||||
|
||||
def init_app(self, app: Flask):
|
||||
"""Initialize cache regions"""
|
||||
self._regions = create_cache_regions(app)
|
||||
|
||||
# Store regions in instance
|
||||
for region_name, region in self._regions.items():
|
||||
setattr(self, f"{region_name}_region", region)
|
||||
|
||||
app.logger.info(f'Cache regions initialized: {self._regions.keys()}')
|
||||
|
||||
def register_handler(self, handler_class: Type[CacheHandler], region: str):
|
||||
"""Register a cache handler class with its region"""
|
||||
if not hasattr(handler_class, 'handler_name'):
|
||||
raise ValueError("Cache handler must define handler_name class attribute")
|
||||
self._handlers[handler_class] = region
|
||||
|
||||
# Create handler instance
|
||||
region_instance = self._regions[region]
|
||||
handler_instance = handler_class(region_instance)
|
||||
self._handler_instances[handler_class.handler_name] = handler_instance
|
||||
|
||||
def invalidate_region(self, region_name: str):
|
||||
"""Invalidate an entire cache region"""
|
||||
if region_name in self._regions:
|
||||
self._regions[region_name].invalidate()
|
||||
else:
|
||||
raise ValueError(f"Unknown cache region: {region_name}")
|
||||
|
||||
def __getattr__(self, name):
|
||||
"""Handle dynamic access to registered handlers"""
|
||||
instances = object.__getattribute__(self, '_handler_instances')
|
||||
if name in instances:
|
||||
return instances[name]
|
||||
raise AttributeError(f"'EveAICacheManager' object has no attribute '{name}'")
|
||||
102
common/utils/cache/license_cache.py
vendored
Normal file
102
common/utils/cache/license_cache.py
vendored
Normal file
@@ -0,0 +1,102 @@
|
||||
# common/utils/cache/license_cache.py
|
||||
from typing import Dict, Any, Optional
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
|
||||
from flask import current_app
|
||||
from sqlalchemy import and_
|
||||
from sqlalchemy.inspection import inspect
|
||||
|
||||
from common.utils.cache.base import CacheHandler
|
||||
from common.models.entitlements import License
|
||||
|
||||
|
||||
class LicenseCacheHandler(CacheHandler[License]):
|
||||
"""Handles caching of active licenses for tenants"""
|
||||
handler_name = 'license_cache'
|
||||
|
||||
def __init__(self, region):
|
||||
super().__init__(region, 'active_license')
|
||||
self.configure_keys('tenant_id')
|
||||
|
||||
def _to_cache_data(self, instance: License) -> Dict[str, Any]:
|
||||
"""Convert License instance to cache data using SQLAlchemy inspection"""
|
||||
if not instance:
|
||||
return {}
|
||||
|
||||
# Get all column attributes from the SQLAlchemy model
|
||||
mapper = inspect(License)
|
||||
data = {}
|
||||
|
||||
for column in mapper.columns:
|
||||
value = getattr(instance, column.name)
|
||||
|
||||
# Handle date serialization
|
||||
if isinstance(value, dt):
|
||||
data[column.name] = value.isoformat()
|
||||
else:
|
||||
data[column.name] = value
|
||||
|
||||
return data
|
||||
|
||||
def _from_cache_data(self, data: Dict[str, Any], **kwargs) -> License:
|
||||
"""Create License instance from cache data using SQLAlchemy inspection"""
|
||||
if not data:
|
||||
return None
|
||||
|
||||
# Create a new License instance
|
||||
license = License()
|
||||
mapper = inspect(License)
|
||||
|
||||
# Set all attributes dynamically
|
||||
for column in mapper.columns:
|
||||
if column.name in data:
|
||||
value = data[column.name]
|
||||
|
||||
# Handle date deserialization
|
||||
if column.name.endswith('_date') and value:
|
||||
if isinstance(value, str):
|
||||
value = dt.fromisoformat(value).date()
|
||||
|
||||
setattr(license, column.name, value)
|
||||
|
||||
return license
|
||||
|
||||
def _should_cache(self, value: License) -> bool:
|
||||
"""Validate if the license should be cached"""
|
||||
return value is not None and value.id is not None
|
||||
|
||||
def get_active_license(self, tenant_id: int) -> Optional[License]:
|
||||
"""
|
||||
Get the currently active license for a tenant
|
||||
|
||||
Args:
|
||||
tenant_id: ID of the tenant
|
||||
|
||||
Returns:
|
||||
License instance if found, None otherwise
|
||||
"""
|
||||
|
||||
def creator_func(tenant_id: int) -> Optional[License]:
|
||||
from common.extensions import db
|
||||
current_date = dt.now(tz=tz.utc).date()
|
||||
|
||||
# TODO --> Active License via active Period?
|
||||
|
||||
return (db.session.query(License)
|
||||
.filter_by(tenant_id=tenant_id)
|
||||
.filter(License.start_date <= current_date)
|
||||
.last())
|
||||
|
||||
return self.get(creator_func, tenant_id=tenant_id)
|
||||
|
||||
def invalidate_tenant_license(self, tenant_id: int):
|
||||
"""Invalidate cached license for specific tenant"""
|
||||
self.invalidate(tenant_id=tenant_id)
|
||||
|
||||
|
||||
def register_license_cache_handlers(cache_manager) -> None:
|
||||
"""Register license cache handlers with cache manager"""
|
||||
cache_manager.register_handler(
|
||||
LicenseCacheHandler,
|
||||
'eveai_model' # Use existing eveai_model region
|
||||
)
|
||||
74
common/utils/cache/regions.py
vendored
Normal file
74
common/utils/cache/regions.py
vendored
Normal file
@@ -0,0 +1,74 @@
|
||||
# common/utils/cache/regions.py
|
||||
import time
|
||||
|
||||
from dogpile.cache import make_region
|
||||
from urllib.parse import urlparse
|
||||
import os
|
||||
|
||||
|
||||
def get_redis_config(app):
|
||||
"""
|
||||
Create Redis configuration dict based on app config
|
||||
Handles both authenticated and non-authenticated setups
|
||||
"""
|
||||
# Parse the REDIS_BASE_URI to get all components
|
||||
redis_uri = urlparse(app.config['REDIS_BASE_URI'])
|
||||
|
||||
config = {
|
||||
'host': redis_uri.hostname,
|
||||
'port': int(redis_uri.port or 6379),
|
||||
'db': 4, # Keep this for later use
|
||||
'redis_expiration_time': 3600,
|
||||
'distributed_lock': True,
|
||||
'thread_local_lock': False,
|
||||
}
|
||||
|
||||
# Add authentication if provided
|
||||
if redis_uri.username and redis_uri.password:
|
||||
config.update({
|
||||
'username': redis_uri.username,
|
||||
'password': redis_uri.password
|
||||
})
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def create_cache_regions(app):
|
||||
"""Initialize all cache regions with app config"""
|
||||
redis_config = get_redis_config(app)
|
||||
regions = {}
|
||||
startup_time = int(time.time())
|
||||
|
||||
# Region for model-related caching (ModelVariables etc)
|
||||
model_region = make_region(name='eveai_model').configure(
|
||||
'dogpile.cache.redis',
|
||||
arguments=redis_config,
|
||||
replace_existing_backend=True
|
||||
)
|
||||
regions['eveai_model'] = model_region
|
||||
|
||||
# Region for eveai_chat_workers components (Specialists, Retrievers, ...)
|
||||
eveai_chat_workers_region = make_region(name='eveai_chat_workers').configure(
|
||||
'dogpile.cache.redis',
|
||||
arguments=redis_config, # arguments={**redis_config, 'db': 4}, # Different DB
|
||||
replace_existing_backend=True
|
||||
)
|
||||
regions['eveai_chat_workers'] = eveai_chat_workers_region
|
||||
|
||||
# Region for eveai_workers components (Processors, ...)
|
||||
eveai_workers_region = make_region(name='eveai_workers').configure(
|
||||
'dogpile.cache.redis',
|
||||
arguments=redis_config, # Same config for now
|
||||
replace_existing_backend=True
|
||||
)
|
||||
regions['eveai_workers'] = eveai_workers_region
|
||||
|
||||
eveai_config_region = make_region(name='eveai_config').configure(
|
||||
'dogpile.cache.redis',
|
||||
arguments=redis_config,
|
||||
replace_existing_backend=True
|
||||
)
|
||||
regions['eveai_config'] = eveai_config_region
|
||||
|
||||
return regions
|
||||
|
||||
@@ -8,8 +8,6 @@ celery_app = Celery()
|
||||
|
||||
def init_celery(celery, app, is_beat=False):
|
||||
celery_app.main = app.name
|
||||
app.logger.debug(f'CELERY_BROKER_URL: {app.config["CELERY_BROKER_URL"]}')
|
||||
app.logger.debug(f'CELERY_RESULT_BACKEND: {app.config["CELERY_RESULT_BACKEND"]}')
|
||||
|
||||
celery_config = {
|
||||
'broker_url': app.config.get('CELERY_BROKER_URL', 'redis://localhost:6379/0'),
|
||||
@@ -60,39 +58,6 @@ def init_celery(celery, app, is_beat=False):
|
||||
|
||||
celery.Task = ContextTask
|
||||
|
||||
# Original init_celery before updating for beat
|
||||
# def init_celery(celery, app):
|
||||
# celery_app.main = app.name
|
||||
# app.logger.debug(f'CELERY_BROKER_URL: {app.config["CELERY_BROKER_URL"]}')
|
||||
# app.logger.debug(f'CELERY_RESULT_BACKEND: {app.config["CELERY_RESULT_BACKEND"]}')
|
||||
# celery_config = {
|
||||
# 'broker_url': app.config.get('CELERY_BROKER_URL', 'redis://localhost:6379/0'),
|
||||
# 'result_backend': app.config.get('CELERY_RESULT_BACKEND', 'redis://localhost:6379/0'),
|
||||
# 'task_serializer': app.config.get('CELERY_TASK_SERIALIZER', 'json'),
|
||||
# 'result_serializer': app.config.get('CELERY_RESULT_SERIALIZER', 'json'),
|
||||
# 'accept_content': app.config.get('CELERY_ACCEPT_CONTENT', ['json']),
|
||||
# 'timezone': app.config.get('CELERY_TIMEZONE', 'UTC'),
|
||||
# 'enable_utc': app.config.get('CELERY_ENABLE_UTC', True),
|
||||
# 'task_routes': {'eveai_worker.tasks.create_embeddings': {'queue': 'embeddings',
|
||||
# 'routing_key': 'embeddings.create_embeddings'}},
|
||||
# }
|
||||
# celery_app.conf.update(**celery_config)
|
||||
#
|
||||
# # Setting up Celery task queues
|
||||
# celery_app.conf.task_queues = (
|
||||
# Queue('default', routing_key='task.#'),
|
||||
# Queue('embeddings', routing_key='embeddings.#', queue_arguments={'x-max-priority': 10}),
|
||||
# Queue('llm_interactions', routing_key='llm_interactions.#', queue_arguments={'x-max-priority': 5}),
|
||||
# )
|
||||
#
|
||||
# # Ensuring tasks execute with Flask application context
|
||||
# class ContextTask(celery.Task):
|
||||
# def __call__(self, *args, **kwargs):
|
||||
# with app.app_context():
|
||||
# return self.run(*args, **kwargs)
|
||||
#
|
||||
# celery.Task = ContextTask
|
||||
|
||||
|
||||
def make_celery(app_name, config):
|
||||
return celery_app
|
||||
|
||||
710
common/utils/config_field_types.py
Normal file
710
common/utils/config_field_types.py
Normal file
@@ -0,0 +1,710 @@
|
||||
from typing import Optional, List, Union, Dict, Any, Pattern
|
||||
from pydantic import BaseModel, field_validator, model_validator
|
||||
from typing_extensions import Annotated
|
||||
import re
|
||||
from datetime import datetime
|
||||
import json
|
||||
from textwrap import dedent
|
||||
import yaml
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
class TaggingField(BaseModel):
|
||||
"""Represents a single tagging field configuration"""
|
||||
type: str
|
||||
required: bool = False
|
||||
description: Optional[str] = None
|
||||
allowed_values: Optional[List[Any]] = None # for enum type
|
||||
min_value: Optional[Union[int, float]] = None # for numeric types
|
||||
max_value: Optional[Union[int, float]] = None # for numeric types
|
||||
|
||||
@field_validator('type', mode='before')
|
||||
@classmethod
|
||||
def validate_type(cls, v: str) -> str:
|
||||
valid_types = ['string', 'integer', 'float', 'date', 'enum']
|
||||
if v not in valid_types:
|
||||
raise ValueError(f'type must be one of {valid_types}')
|
||||
return v
|
||||
|
||||
@model_validator(mode='after')
|
||||
def validate_field_constraints(self) -> 'TaggingField':
|
||||
# Validate enum constraints
|
||||
if self.type == 'enum':
|
||||
if not self.allowed_values:
|
||||
raise ValueError('allowed_values must be provided for enum type')
|
||||
elif self.allowed_values is not None:
|
||||
raise ValueError('allowed_values only valid for enum type')
|
||||
|
||||
# Validate numeric constraints
|
||||
if self.type not in ('integer', 'float'):
|
||||
if self.min_value is not None or self.max_value is not None:
|
||||
raise ValueError('min_value/max_value only valid for numeric types')
|
||||
else:
|
||||
if self.min_value is not None and self.max_value is not None and self.min_value >= self.max_value:
|
||||
raise ValueError('min_value must be less than max_value')
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class TaggingFields(BaseModel):
|
||||
"""Represents a collection of tagging fields, mapped by their names"""
|
||||
fields: Dict[str, TaggingField]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Dict[str, Any]]) -> 'TaggingFields':
|
||||
return cls(fields={
|
||||
field_name: TaggingField(**field_config)
|
||||
for field_name, field_config in data.items()
|
||||
})
|
||||
|
||||
def to_dict(self) -> Dict[str, Dict[str, Any]]:
|
||||
return {
|
||||
field_name: field.model_dump(exclude_none=True)
|
||||
for field_name, field in self.fields.items()
|
||||
}
|
||||
|
||||
|
||||
class ChunkingPatternsField(BaseModel):
|
||||
"""Represents a set of chunking patterns"""
|
||||
patterns: List[str]
|
||||
|
||||
@field_validator('patterns')
|
||||
def validate_patterns(cls, patterns):
|
||||
for pattern in patterns:
|
||||
try:
|
||||
re.compile(pattern)
|
||||
except re.error as e:
|
||||
raise ValueError(f"Invalid regex pattern '{pattern}': {str(e)}")
|
||||
return patterns
|
||||
|
||||
|
||||
class ArgumentConstraint(BaseModel):
|
||||
"""Base class for all argument constraints"""
|
||||
description: Optional[str] = None
|
||||
error_message: Optional[str] = None
|
||||
|
||||
|
||||
class NumericConstraint(ArgumentConstraint):
|
||||
"""Constraints for numeric values (int/float)"""
|
||||
min_value: Optional[float] = None
|
||||
max_value: Optional[float] = None
|
||||
include_min: bool = True # True for >= min_value, False for > min_value
|
||||
include_max: bool = True # True for <= max_value, False for < max_value
|
||||
|
||||
@model_validator(mode='after')
|
||||
def validate_ranges(self) -> 'NumericConstraint':
|
||||
if self.min_value is not None and self.max_value is not None:
|
||||
if self.min_value > self.max_value:
|
||||
raise ValueError("min_value must be less than or equal to max_value")
|
||||
return self
|
||||
|
||||
def validate(self, value: Union[int, float]) -> bool:
|
||||
if self.min_value is not None:
|
||||
if self.include_min and value < self.min_value:
|
||||
return False
|
||||
if not self.include_min and value <= self.min_value:
|
||||
return False
|
||||
if self.max_value is not None:
|
||||
if self.include_max and value > self.max_value:
|
||||
return False
|
||||
if not self.include_max and value >= self.max_value:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class StringConstraint(ArgumentConstraint):
|
||||
"""Constraints for string values"""
|
||||
min_length: Optional[int] = None
|
||||
max_length: Optional[int] = None
|
||||
patterns: Optional[List[str]] = None # List of regex patterns to match
|
||||
pattern_match_all: bool = False # If True, string must match all patterns
|
||||
forbidden_patterns: Optional[List[str]] = None # List of regex patterns that must not match
|
||||
allow_empty: bool = False
|
||||
|
||||
@field_validator('patterns', 'forbidden_patterns')
|
||||
@classmethod
|
||||
def validate_patterns(cls, v: Optional[List[str]]) -> Optional[List[str]]:
|
||||
if v is not None:
|
||||
# Validate each pattern compiles
|
||||
for pattern in v:
|
||||
try:
|
||||
re.compile(pattern)
|
||||
except re.error as e:
|
||||
raise ValueError(f"Invalid regex pattern '{pattern}': {str(e)}")
|
||||
return v
|
||||
|
||||
def validate(self, value: str) -> bool:
|
||||
if not self.allow_empty and not value:
|
||||
return False
|
||||
|
||||
if self.min_length is not None and len(value) < self.min_length:
|
||||
return False
|
||||
|
||||
if self.max_length is not None and len(value) > self.max_length:
|
||||
return False
|
||||
|
||||
if self.patterns:
|
||||
matches = [bool(re.search(pattern, value)) for pattern in self.patterns]
|
||||
if self.pattern_match_all and not all(matches):
|
||||
return False
|
||||
if not self.pattern_match_all and not any(matches):
|
||||
return False
|
||||
|
||||
if self.forbidden_patterns:
|
||||
for pattern in self.forbidden_patterns:
|
||||
if re.search(pattern, value):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class DateConstraint(ArgumentConstraint):
|
||||
"""Constraints for date values"""
|
||||
min_date: Optional[datetime] = None
|
||||
max_date: Optional[datetime] = None
|
||||
include_min: bool = True
|
||||
include_max: bool = True
|
||||
allowed_formats: Optional[List[str]] = None # List of allowed date formats
|
||||
|
||||
@model_validator(mode='after')
|
||||
def validate_ranges(self) -> 'DateConstraint':
|
||||
if self.min_date and self.max_date and self.min_date > self.max_date:
|
||||
raise ValueError("min_date must be less than or equal to max_date")
|
||||
return self
|
||||
|
||||
def validate(self, value: datetime) -> bool:
|
||||
if self.min_date is not None:
|
||||
if self.include_min and value < self.min_date:
|
||||
return False
|
||||
if not self.include_min and value <= self.min_date:
|
||||
return False
|
||||
|
||||
if self.max_date is not None:
|
||||
if self.include_max and value > self.max_date:
|
||||
return False
|
||||
if not self.include_max and value >= self.max_date:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class EnumConstraint(ArgumentConstraint):
|
||||
"""Constraints for enum values"""
|
||||
allowed_values: List[Any]
|
||||
case_sensitive: bool = True # For string enums
|
||||
allow_multiple: bool = False # If True, value can be a list of allowed values
|
||||
min_selections: Optional[int] = None # When allow_multiple is True
|
||||
max_selections: Optional[int] = None # When allow_multiple is True
|
||||
|
||||
@model_validator(mode='after')
|
||||
def validate_selections(self) -> 'EnumConstraint':
|
||||
if self.allow_multiple:
|
||||
if self.min_selections is not None and self.max_selections is not None:
|
||||
if self.min_selections > self.max_selections:
|
||||
raise ValueError("min_selections must be less than or equal to max_selections")
|
||||
if self.max_selections > len(self.allowed_values):
|
||||
raise ValueError("max_selections cannot be greater than number of allowed values")
|
||||
return self
|
||||
|
||||
def validate(self, value: Union[Any, List[Any]]) -> bool:
|
||||
if self.allow_multiple:
|
||||
if not isinstance(value, list):
|
||||
return False
|
||||
|
||||
if self.min_selections is not None and len(value) < self.min_selections:
|
||||
return False
|
||||
|
||||
if self.max_selections is not None and len(value) > self.max_selections:
|
||||
return False
|
||||
|
||||
for v in value:
|
||||
if not self._validate_single_value(v):
|
||||
return False
|
||||
else:
|
||||
return self._validate_single_value(value)
|
||||
|
||||
return True
|
||||
|
||||
def _validate_single_value(self, value: Any) -> bool:
|
||||
if isinstance(value, str) and not self.case_sensitive:
|
||||
return any(str(value).lower() == str(v).lower() for v in self.allowed_values)
|
||||
return value in self.allowed_values
|
||||
|
||||
|
||||
class ArgumentDefinition(BaseModel):
|
||||
"""Defines an argument with its type and constraints"""
|
||||
name: str
|
||||
type: str
|
||||
description: Optional[str] = None
|
||||
required: bool = False
|
||||
default: Optional[Any] = None
|
||||
constraints: Optional[Union[NumericConstraint, StringConstraint, DateConstraint, EnumConstraint]] = None
|
||||
|
||||
@field_validator('type')
|
||||
@classmethod
|
||||
def validate_type(cls, v: str) -> str:
|
||||
valid_types = ['string', 'integer', 'float', 'date', 'enum']
|
||||
if v not in valid_types:
|
||||
raise ValueError(f'type must be one of {valid_types}')
|
||||
return v
|
||||
|
||||
@model_validator(mode='after')
|
||||
def validate_constraints(self) -> 'ArgumentDefinition':
|
||||
if self.constraints:
|
||||
expected_constraint_types = {
|
||||
'string': StringConstraint,
|
||||
'integer': NumericConstraint,
|
||||
'float': NumericConstraint,
|
||||
'date': DateConstraint,
|
||||
'enum': EnumConstraint
|
||||
}
|
||||
|
||||
expected_type = expected_constraint_types.get(self.type)
|
||||
if not isinstance(self.constraints, expected_type):
|
||||
raise ValueError(f'Constraints for type {self.type} must be of type {expected_type.__name__}')
|
||||
|
||||
if self.default is not None:
|
||||
if not self.constraints.validate(self.default):
|
||||
raise ValueError(f'Default value does not satisfy constraints for {self.name}')
|
||||
|
||||
return self
|
||||
|
||||
|
||||
class ArgumentDefinitions(BaseModel):
|
||||
"""Collection of argument definitions"""
|
||||
arguments: Dict[str, ArgumentDefinition]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Dict[str, Any]]) -> 'ArgumentDefinitions':
|
||||
return cls(arguments={
|
||||
arg_name: ArgumentDefinition(**arg_config)
|
||||
for arg_name, arg_config in data.items()
|
||||
})
|
||||
|
||||
def to_dict(self) -> Dict[str, Dict[str, Any]]:
|
||||
return {
|
||||
arg_name: arg.model_dump(exclude_none=True)
|
||||
for arg_name, arg in self.arguments.items()
|
||||
}
|
||||
|
||||
def validate_argument_values(self, values: Dict[str, Any]) -> Dict[str, str]:
|
||||
"""
|
||||
Validate a set of argument values against their definitions
|
||||
Returns a dictionary of error messages for invalid arguments
|
||||
"""
|
||||
errors = {}
|
||||
|
||||
# Check for required arguments
|
||||
for name, arg_def in self.arguments.items():
|
||||
if arg_def.required and name not in values:
|
||||
errors[name] = "Required argument missing"
|
||||
continue
|
||||
|
||||
if name in values:
|
||||
value = values[name]
|
||||
|
||||
# Validate type
|
||||
try:
|
||||
if arg_def.type == 'integer':
|
||||
value = int(value)
|
||||
elif arg_def.type == 'float':
|
||||
value = float(value)
|
||||
elif arg_def.type == 'date' and isinstance(value, str):
|
||||
if arg_def.constraints and arg_def.constraints.allowed_formats:
|
||||
for fmt in arg_def.constraints.allowed_formats:
|
||||
try:
|
||||
value = datetime.strptime(value, fmt)
|
||||
break
|
||||
except ValueError:
|
||||
continue
|
||||
else:
|
||||
errors[
|
||||
name] = f"Invalid date format. Allowed formats: {arg_def.constraints.allowed_formats}"
|
||||
continue
|
||||
except (ValueError, TypeError):
|
||||
errors[name] = f"Invalid type. Expected {arg_def.type}"
|
||||
continue
|
||||
|
||||
# Validate constraints
|
||||
if arg_def.constraints and not arg_def.constraints.validate(value):
|
||||
errors[name] = arg_def.constraints.error_message or "Value does not satisfy constraints"
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentationFormat:
|
||||
"""Constants for documentation formats"""
|
||||
MARKDOWN = "markdown"
|
||||
JSON = "json"
|
||||
YAML = "yaml"
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentationVersion:
|
||||
"""Constants for documentation versions"""
|
||||
BASIC = "basic" # Original documentation without retriever info
|
||||
EXTENDED = "extended" # Including retriever documentation
|
||||
|
||||
|
||||
def _generate_argument_constraints(field_config: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""Generate possible argument constraints based on field type"""
|
||||
constraints = []
|
||||
|
||||
base_constraint = {
|
||||
"description": f"Constraint for {field_config.get('description', 'field')}",
|
||||
"error_message": "Optional custom error message"
|
||||
}
|
||||
|
||||
if field_config["type"] == "integer" or field_config["type"] == "float":
|
||||
constraints.append({
|
||||
**base_constraint,
|
||||
"type": "NumericConstraint",
|
||||
"possible_constraints": {
|
||||
"min_value": "number",
|
||||
"max_value": "number",
|
||||
"include_min": "boolean",
|
||||
"include_max": "boolean"
|
||||
},
|
||||
"example": {
|
||||
"min_value": field_config.get("min_value", 0),
|
||||
"max_value": field_config.get("max_value", 100),
|
||||
"include_min": True,
|
||||
"include_max": True
|
||||
}
|
||||
})
|
||||
|
||||
elif field_config["type"] == "string":
|
||||
constraints.append({
|
||||
**base_constraint,
|
||||
"type": "StringConstraint",
|
||||
"possible_constraints": {
|
||||
"min_length": "integer",
|
||||
"max_length": "integer",
|
||||
"patterns": "list[str]",
|
||||
"pattern_match_all": "boolean",
|
||||
"forbidden_patterns": "list[str]",
|
||||
"allow_empty": "boolean"
|
||||
},
|
||||
"example": {
|
||||
"min_length": 1,
|
||||
"max_length": 100,
|
||||
"patterns": ["^[A-Za-z0-9]+$"],
|
||||
"pattern_match_all": False,
|
||||
"forbidden_patterns": ["^test_", "_temp$"],
|
||||
"allow_empty": False
|
||||
}
|
||||
})
|
||||
|
||||
elif field_config["type"] == "enum":
|
||||
constraints.append({
|
||||
**base_constraint,
|
||||
"type": "EnumConstraint",
|
||||
"possible_constraints": {
|
||||
"allowed_values": f"list[{field_config.get('allowed_values', ['value1', 'value2'])}]",
|
||||
"case_sensitive": "boolean",
|
||||
"allow_multiple": "boolean",
|
||||
"min_selections": "integer",
|
||||
"max_selections": "integer"
|
||||
},
|
||||
"example": {
|
||||
"allowed_values": field_config.get("allowed_values", ["value1", "value2"]),
|
||||
"case_sensitive": True,
|
||||
"allow_multiple": True,
|
||||
"min_selections": 1,
|
||||
"max_selections": 2
|
||||
}
|
||||
})
|
||||
|
||||
elif field_config["type"] == "date":
|
||||
constraints.append({
|
||||
**base_constraint,
|
||||
"type": "DateConstraint",
|
||||
"possible_constraints": {
|
||||
"min_date": "datetime",
|
||||
"max_date": "datetime",
|
||||
"include_min": "boolean",
|
||||
"include_max": "boolean",
|
||||
"allowed_formats": "list[str]"
|
||||
},
|
||||
"example": {
|
||||
"min_date": "2024-01-01T00:00:00",
|
||||
"max_date": "2024-12-31T23:59:59",
|
||||
"include_min": True,
|
||||
"include_max": True,
|
||||
"allowed_formats": ["%Y-%m-%d", "%Y/%m/%d"]
|
||||
}
|
||||
})
|
||||
|
||||
return constraints
|
||||
|
||||
|
||||
def generate_field_documentation(
|
||||
tagging_fields: Dict[str, Any],
|
||||
format: str = "markdown",
|
||||
version: str = "basic"
|
||||
) -> str:
|
||||
"""
|
||||
Generate documentation for tagging fields configuration.
|
||||
|
||||
Args:
|
||||
tagging_fields: Dictionary containing tagging fields configuration
|
||||
format: Output format ("markdown", "json", or "yaml")
|
||||
version: Documentation version ("basic" or "extended")
|
||||
|
||||
Returns:
|
||||
str: Formatted documentation
|
||||
"""
|
||||
if version not in [DocumentationVersion.BASIC, DocumentationVersion.EXTENDED]:
|
||||
raise ValueError(f"Unsupported documentation version: {version}")
|
||||
|
||||
# Normalize fields configuration
|
||||
normalized_fields = {}
|
||||
|
||||
for field_name, field_config in tagging_fields.items():
|
||||
field_doc = {
|
||||
"name": field_name,
|
||||
"type": field_config["type"],
|
||||
"required": field_config.get("required", False),
|
||||
"description": field_config.get("description", "No description provided"),
|
||||
"constraints": []
|
||||
}
|
||||
|
||||
# Only include possible arguments in extended version
|
||||
if version == DocumentationVersion.EXTENDED:
|
||||
field_doc["possible_arguments"] = _generate_argument_constraints(field_config)
|
||||
|
||||
# Add type-specific constraints
|
||||
if field_config["type"] == "integer" or field_config["type"] == "float":
|
||||
if "min_value" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Minimum value: {field_config['min_value']}")
|
||||
if "max_value" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Maximum value: {field_config['max_value']}")
|
||||
|
||||
elif field_config["type"] == "string":
|
||||
if "min_length" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Minimum length: {field_config['min_length']}")
|
||||
if "max_length" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Maximum length: {field_config['max_length']}")
|
||||
if "patterns" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Must match patterns: {', '.join(field_config['patterns'])}")
|
||||
|
||||
elif field_config["type"] == "enum":
|
||||
if "allowed_values" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Allowed values: {', '.join(str(v) for v in field_config['allowed_values'])}")
|
||||
|
||||
elif field_config["type"] == "date":
|
||||
if "min_date" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Minimum date: {field_config['min_date']}")
|
||||
if "max_date" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Maximum date: {field_config['max_date']}")
|
||||
if "allowed_formats" in field_config:
|
||||
field_doc["constraints"].append(
|
||||
f"Allowed formats: {', '.join(field_config['allowed_formats'])}")
|
||||
|
||||
normalized_fields[field_name] = field_doc
|
||||
|
||||
# Generate documentation in requested format
|
||||
if format == DocumentationFormat.MARKDOWN:
|
||||
return _generate_markdown_docs(normalized_fields, version)
|
||||
elif format == DocumentationFormat.JSON:
|
||||
return _generate_json_docs(normalized_fields, version)
|
||||
elif format == DocumentationFormat.YAML:
|
||||
return _generate_yaml_docs(normalized_fields, version)
|
||||
else:
|
||||
raise ValueError(f"Unsupported documentation format: {format}")
|
||||
|
||||
|
||||
def _generate_markdown_docs(fields: Dict[str, Any], version: str) -> str:
|
||||
"""Generate markdown documentation"""
|
||||
docs = ["# Tagging Fields Documentation\n"]
|
||||
|
||||
# Add overview table
|
||||
docs.append("## Fields Overview\n")
|
||||
docs.append("| Field Name | Type | Required | Description |")
|
||||
docs.append("|------------|------|----------|-------------|")
|
||||
|
||||
for field_name, field in fields.items():
|
||||
docs.append(
|
||||
f"| {field_name} | {field['type']} | "
|
||||
f"{'Yes' if field['required'] else 'No'} | {field['description']} |"
|
||||
)
|
||||
|
||||
# Add detailed field specifications
|
||||
docs.append("\n## Detailed Field Specifications\n")
|
||||
|
||||
for field_name, field in fields.items():
|
||||
docs.append(f"### {field_name}\n")
|
||||
docs.append(f"**Type:** {field['type']}")
|
||||
docs.append(f"**Required:** {'Yes' if field['required'] else 'No'}")
|
||||
docs.append(f"**Description:** {field['description']}\n")
|
||||
|
||||
if field["constraints"]:
|
||||
docs.append("**Field Constraints:**")
|
||||
for constraint in field["constraints"]:
|
||||
docs.append(f"- {constraint}")
|
||||
docs.append("")
|
||||
|
||||
# Add retriever argument documentation only in extended version
|
||||
if version == DocumentationVersion.EXTENDED and "possible_arguments" in field:
|
||||
docs.append("**Possible Retriever Arguments:**")
|
||||
for arg_constraint in field["possible_arguments"]:
|
||||
docs.append(f"\n*{arg_constraint['type']}*")
|
||||
docs.append(f"Description: {arg_constraint['description']}")
|
||||
docs.append("\nPossible constraints:")
|
||||
for const_name, const_type in arg_constraint["possible_constraints"].items():
|
||||
docs.append(f"- `{const_name}`: {const_type}")
|
||||
|
||||
docs.append("\nExample:")
|
||||
docs.append("```python")
|
||||
docs.append(json.dumps(arg_constraint["example"], indent=2))
|
||||
docs.append("```\n")
|
||||
|
||||
# Add example retriever configuration only in extended version
|
||||
if version == DocumentationVersion.EXTENDED:
|
||||
docs.append("\n## Example Retriever Configuration\n")
|
||||
docs.append("```python")
|
||||
example_config = {
|
||||
"metadata_filters": {
|
||||
field_name: field["possible_arguments"][0]["example"]
|
||||
for field_name, field in fields.items()
|
||||
if "possible_arguments" in field
|
||||
}
|
||||
}
|
||||
docs.append(json.dumps(example_config, indent=2))
|
||||
docs.append("```")
|
||||
|
||||
return "\n".join(docs)
|
||||
|
||||
|
||||
def _generate_json_docs(fields: Dict[str, Any], version: str) -> str:
|
||||
"""Generate JSON documentation"""
|
||||
doc = {
|
||||
"tagging_fields_documentation": {
|
||||
"version": version,
|
||||
"fields": fields
|
||||
}
|
||||
}
|
||||
|
||||
if version == DocumentationVersion.EXTENDED:
|
||||
doc["tagging_fields_documentation"]["example_retriever_config"] = {
|
||||
"metadata_filters": {
|
||||
field_name: field["possible_arguments"][0]["example"]
|
||||
for field_name, field in fields.items()
|
||||
if "possible_arguments" in field
|
||||
}
|
||||
}
|
||||
|
||||
return json.dumps(doc, indent=2)
|
||||
|
||||
|
||||
def _generate_yaml_docs(fields: Dict[str, Any], version: str) -> str:
|
||||
"""Generate YAML documentation"""
|
||||
doc = {
|
||||
"tagging_fields_documentation": {
|
||||
"version": version,
|
||||
"fields": fields
|
||||
}
|
||||
}
|
||||
|
||||
if version == DocumentationVersion.EXTENDED:
|
||||
doc["tagging_fields_documentation"]["example_retriever_config"] = {
|
||||
"metadata_filters": {
|
||||
field_name: field["possible_arguments"][0]["example"]
|
||||
for field_name, field in fields.items()
|
||||
if "possible_arguments" in field
|
||||
}
|
||||
}
|
||||
|
||||
return yaml.dump(doc, sort_keys=False, default_flow_style=False)
|
||||
|
||||
|
||||
def patterns_to_json(text_area_content: str) -> str:
|
||||
"""Convert line-based patterns to JSON"""
|
||||
text_area_content = text_area_content.strip()
|
||||
if len(text_area_content) == 0:
|
||||
return json.dumps([])
|
||||
# Split on newlines and remove empty lines
|
||||
patterns = [line.strip() for line in text_area_content.split('\n') if line.strip()]
|
||||
return json.dumps(patterns)
|
||||
|
||||
|
||||
def json_to_patterns(json_content: str) -> str:
|
||||
"""Convert JSON patterns list to text area content"""
|
||||
try:
|
||||
patterns = json.loads(json_content)
|
||||
if not isinstance(patterns, list):
|
||||
raise ValueError("JSON must contain a list of patterns")
|
||||
# Join with newlines
|
||||
return '\n'.join(patterns)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid JSON format: {e}")
|
||||
|
||||
|
||||
def json_to_pattern_list(json_content: str) -> list:
|
||||
"""Convert JSON patterns list to text area content"""
|
||||
try:
|
||||
if json_content:
|
||||
patterns = json.loads(json_content)
|
||||
if not isinstance(patterns, list):
|
||||
raise ValueError("JSON must contain a list of patterns")
|
||||
# Unescape if needed
|
||||
patterns = [pattern.replace('\\\\', '\\') for pattern in patterns]
|
||||
return patterns
|
||||
else:
|
||||
return []
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Invalid JSON format: {e}")
|
||||
|
||||
|
||||
def normalize_json_field(value: str | dict | None, field_name: str = "JSON field") -> dict:
|
||||
"""
|
||||
Normalize a JSON field value to ensure it's a valid dictionary.
|
||||
|
||||
Args:
|
||||
value: The input value which can be:
|
||||
- None (will return empty dict)
|
||||
- String (will be parsed as JSON)
|
||||
- Dict (will be validated and returned)
|
||||
field_name: Name of the field for error messages
|
||||
|
||||
Returns:
|
||||
dict: The normalized JSON data as a Python dictionary
|
||||
|
||||
Raises:
|
||||
ValueError: If the input string is not valid JSON or the input dict contains invalid types
|
||||
"""
|
||||
# Handle None case
|
||||
if value is None:
|
||||
return {}
|
||||
|
||||
# Handle dictionary case
|
||||
if isinstance(value, dict):
|
||||
try:
|
||||
# Validate all values are JSON serializable
|
||||
import json
|
||||
json.dumps(value)
|
||||
return value
|
||||
except TypeError as e:
|
||||
raise ValueError(f"{field_name} contains invalid types: {str(e)}")
|
||||
|
||||
# Handle string case
|
||||
if isinstance(value, str):
|
||||
if not value.strip():
|
||||
return {}
|
||||
|
||||
try:
|
||||
import json
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"{field_name} contains invalid JSON: {str(e)}")
|
||||
|
||||
raise ValueError(f"{field_name} must be a string, dictionary, or None (got {type(value)})")
|
||||
@@ -1,14 +1,14 @@
|
||||
from flask import request, current_app, session
|
||||
from flask_jwt_extended import decode_token, verify_jwt_in_request, get_jwt_identity
|
||||
|
||||
from common.models.user import Tenant, TenantDomain
|
||||
|
||||
|
||||
def get_allowed_origins(tenant_id):
|
||||
session_key = f"allowed_origins_{tenant_id}"
|
||||
if session_key in session:
|
||||
current_app.logger.debug(f"Fetching allowed origins for tenant {tenant_id} from session")
|
||||
return session[session_key]
|
||||
|
||||
current_app.logger.debug(f"Fetching allowed origins for tenant {tenant_id} from database")
|
||||
tenant_domains = TenantDomain.query.filter_by(tenant_id=int(tenant_id)).all()
|
||||
allowed_origins = [domain.domain for domain in tenant_domains]
|
||||
|
||||
@@ -18,51 +18,52 @@ def get_allowed_origins(tenant_id):
|
||||
|
||||
|
||||
def cors_after_request(response, prefix):
|
||||
current_app.logger.debug(f'CORS after request: {request.path}, prefix: {prefix}')
|
||||
current_app.logger.debug(f'request.headers: {request.headers}')
|
||||
current_app.logger.debug(f'request.args: {request.args}')
|
||||
current_app.logger.debug(f'request is json?: {request.is_json}')
|
||||
|
||||
# Exclude health checks from checks
|
||||
if request.path.startswith('/healthz') or request.path.startswith('/_healthz'):
|
||||
current_app.logger.debug('Skipping CORS headers for health checks')
|
||||
response.headers.add('Access-Control-Allow-Origin', '*')
|
||||
response.headers.add('Access-Control-Allow-Headers', '*')
|
||||
response.headers.add('Access-Control-Allow-Methods', '*')
|
||||
return response
|
||||
|
||||
# Handle OPTIONS preflight requests
|
||||
if request.method == 'OPTIONS':
|
||||
response.headers.add('Access-Control-Allow-Origin', '*')
|
||||
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization,X-Tenant-ID')
|
||||
response.headers.add('Access-Control-Allow-Methods', 'GET,POST,PUT,DELETE,OPTIONS')
|
||||
response.headers.add('Access-Control-Allow-Credentials', 'true')
|
||||
return response
|
||||
|
||||
tenant_id = None
|
||||
allowed_origins = []
|
||||
|
||||
# Try to get tenant_id from JSON payload
|
||||
json_data = request.get_json(silent=True)
|
||||
current_app.logger.debug(f'request.get_json(silent=True): {json_data}')
|
||||
|
||||
if json_data and 'tenant_id' in json_data:
|
||||
tenant_id = json_data['tenant_id']
|
||||
# Check Socket.IO connection
|
||||
if 'socket.io' in request.path:
|
||||
token = request.args.get('token')
|
||||
if token:
|
||||
try:
|
||||
decoded = decode_token(token)
|
||||
tenant_id = decoded['sub']
|
||||
except Exception as e:
|
||||
current_app.logger.error(f'Error decoding token: {e}')
|
||||
return response
|
||||
else:
|
||||
# Fallback to get tenant_id from query parameters or headers if JSON is not available
|
||||
tenant_id = request.args.get('tenant_id') or request.args.get('tenantId') or request.headers.get('X-Tenant-ID')
|
||||
|
||||
current_app.logger.debug(f'Identified tenant_id: {tenant_id}')
|
||||
# Regular API requests
|
||||
try:
|
||||
if verify_jwt_in_request(optional=True):
|
||||
tenant_id = get_jwt_identity()
|
||||
except Exception as e:
|
||||
current_app.logger.error(f'Error verifying JWT: {e}')
|
||||
return response
|
||||
|
||||
if tenant_id:
|
||||
origin = request.headers.get('Origin')
|
||||
allowed_origins = get_allowed_origins(tenant_id)
|
||||
current_app.logger.debug(f'Allowed origins for tenant {tenant_id}: {allowed_origins}')
|
||||
else:
|
||||
current_app.logger.warning('tenant_id not found in request')
|
||||
|
||||
origin = request.headers.get('Origin')
|
||||
current_app.logger.debug(f'Origin: {origin}')
|
||||
|
||||
if origin in allowed_origins:
|
||||
response.headers.add('Access-Control-Allow-Origin', origin)
|
||||
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
|
||||
response.headers.add('Access-Control-Allow-Methods', 'GET,POST,PUT,DELETE,OPTIONS')
|
||||
response.headers.add('Access-Control-Allow-Credentials', 'true')
|
||||
current_app.logger.debug(f'CORS headers set for origin: {origin}')
|
||||
else:
|
||||
current_app.logger.warning(f'Origin {origin} not allowed')
|
||||
if origin in allowed_origins:
|
||||
response.headers.add('Access-Control-Allow-Origin', origin)
|
||||
response.headers.add('Access-Control-Allow-Headers', 'Content-Type,Authorization')
|
||||
response.headers.add('Access-Control-Allow-Methods', 'GET,POST,PUT,DELETE,OPTIONS')
|
||||
response.headers.add('Access-Control-Allow-Credentials', 'true')
|
||||
|
||||
return response
|
||||
|
||||
|
||||
@@ -1,62 +1,104 @@
|
||||
from flask import request, session
|
||||
import time
|
||||
from flask_security import current_user
|
||||
import json
|
||||
|
||||
|
||||
def log_request_middleware(app):
|
||||
# @app.before_request
|
||||
# def log_request_info():
|
||||
# start_time = time.time()
|
||||
# app.logger.debug(f"Request URL: {request.url}")
|
||||
# app.logger.debug(f"Request Method: {request.method}")
|
||||
# app.logger.debug(f"Request Headers: {request.headers}")
|
||||
# app.logger.debug(f"Time taken for logging request info: {time.time() - start_time} seconds")
|
||||
# try:
|
||||
# app.logger.debug(f"Request Body: {request.get_data()}")
|
||||
# except Exception as e:
|
||||
# app.logger.error(f"Error reading request body: {e}")
|
||||
# app.logger.debug(f"Time taken for logging request body: {time.time() - start_time} seconds")
|
||||
|
||||
# @app.before_request
|
||||
# def check_csrf_token():
|
||||
# start_time = time.time()
|
||||
# if request.method == "POST":
|
||||
# csrf_token = request.form.get("csrf_token")
|
||||
# app.logger.debug(f"CSRF Token: {csrf_token}")
|
||||
# app.logger.debug(f"Time taken for logging CSRF token: {time.time() - start_time} seconds")
|
||||
|
||||
# @app.before_request
|
||||
# def log_user_info():
|
||||
# if current_user and current_user.is_authenticated:
|
||||
# app.logger.debug(f"Before: User ID: {current_user.id}")
|
||||
# app.logger.debug(f"Before: User Email: {current_user.email}")
|
||||
# app.logger.debug(f"Before: User Roles: {current_user.roles}")
|
||||
# else:
|
||||
# app.logger.debug("After: No user logged in")
|
||||
|
||||
@app.before_request
|
||||
def log_session_state_before():
|
||||
app.logger.debug(f'Session state before request: {session.items()}')
|
||||
|
||||
# @app.after_request
|
||||
# def log_response_info(response):
|
||||
# start_time = time.time()
|
||||
# app.logger.debug(f"Response Status: {response.status}")
|
||||
# app.logger.debug(f"Response Headers: {response.headers}")
|
||||
#
|
||||
# app.logger.debug(f"Time taken for logging response info: {time.time() - start_time} seconds")
|
||||
# return response
|
||||
|
||||
# @app.after_request
|
||||
# def log_user_after_request(response):
|
||||
# if current_user and current_user.is_authenticated:
|
||||
# app.logger.debug(f"After: User ID: {current_user.id}")
|
||||
# app.logger.debug(f"after: User Email: {current_user.email}")
|
||||
# app.logger.debug(f"After: User Roles: {current_user.roles}")
|
||||
# else:
|
||||
# app.logger.debug("After: No user logged in")
|
||||
pass
|
||||
|
||||
@app.after_request
|
||||
def log_session_state_after(response):
|
||||
app.logger.debug(f'Session state after request: {session.items()}')
|
||||
return response
|
||||
|
||||
|
||||
def register_request_debugger(app):
|
||||
@app.before_request
|
||||
def debug_request_info():
|
||||
"""Log consolidated request information for debugging"""
|
||||
# Skip health check endpoints
|
||||
if request.path.startswith('/_healthz') or request.path.startswith('/healthz'):
|
||||
return
|
||||
|
||||
# Gather all request information in a structured way
|
||||
debug_info = {
|
||||
"basic_info": {
|
||||
"method": request.method,
|
||||
"path": request.path,
|
||||
"content_type": request.content_type,
|
||||
"content_length": request.content_length
|
||||
},
|
||||
"environment": {
|
||||
"remote_addr": request.remote_addr,
|
||||
"user_agent": str(request.user_agent)
|
||||
}
|
||||
}
|
||||
|
||||
# Add headers (excluding sensitive ones)
|
||||
safe_headers = {k: v for k, v in request.headers.items()
|
||||
if k.lower() not in ('authorization', 'cookie', 'x-api-key')}
|
||||
debug_info["headers"] = safe_headers
|
||||
|
||||
# Add authentication info (presence only)
|
||||
auth_header = request.headers.get('Authorization', '')
|
||||
debug_info["auth_info"] = {
|
||||
"has_auth_header": bool(auth_header),
|
||||
"auth_type": auth_header.split(' ')[0] if auth_header else None,
|
||||
"token_length": len(auth_header.split(' ')[1]) if auth_header and len(auth_header.split(' ')) > 1 else 0,
|
||||
"header_format": 'Valid format' if auth_header.startswith('Bearer ') else 'Invalid format',
|
||||
"raw_header": auth_header[:10] + '...' if auth_header else None # Show first 10 chars only
|
||||
}
|
||||
|
||||
# Add request data based on type
|
||||
if request.is_json:
|
||||
try:
|
||||
json_data = request.get_json()
|
||||
if isinstance(json_data, dict):
|
||||
# Remove sensitive fields from logging
|
||||
safe_json = {k: v for k, v in json_data.items()
|
||||
if not any(sensitive in k.lower()
|
||||
for sensitive in ['password', 'token', 'secret', 'key'])}
|
||||
debug_info["request_data"] = {
|
||||
"type": "json",
|
||||
"content": safe_json
|
||||
}
|
||||
except Exception as e:
|
||||
debug_info["request_data"] = {
|
||||
"type": "json",
|
||||
"error": str(e)
|
||||
}
|
||||
elif request.form:
|
||||
safe_form = {k: v for k, v in request.form.items()
|
||||
if not any(sensitive in k.lower()
|
||||
for sensitive in ['password', 'token', 'secret', 'key'])}
|
||||
debug_info["request_data"] = {
|
||||
"type": "form",
|
||||
"content": safe_form
|
||||
}
|
||||
|
||||
# Add file information if present
|
||||
if request.files:
|
||||
debug_info["files"] = {
|
||||
name: {
|
||||
"filename": f.filename,
|
||||
"content_type": f.content_type,
|
||||
"content_length": f.content_length if hasattr(f, 'content_length') else None
|
||||
}
|
||||
for name, f in request.files.items()
|
||||
}
|
||||
|
||||
# Add CORS information if present
|
||||
cors_headers = {
|
||||
"origin": request.headers.get('Origin'),
|
||||
"request_method": request.headers.get('Access-Control-Request-Method'),
|
||||
"request_headers": request.headers.get('Access-Control-Request-Headers')
|
||||
}
|
||||
if any(cors_headers.values()):
|
||||
debug_info["cors"] = {k: v for k, v in cors_headers.items() if v is not None}
|
||||
|
||||
# Format the debug info as a pretty-printed JSON string with indentation
|
||||
formatted_debug_info = json.dumps(debug_info, indent=2, sort_keys=True)
|
||||
|
||||
|
||||
@@ -3,27 +3,59 @@ from datetime import datetime as dt, timezone as tz
|
||||
from sqlalchemy import desc
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
from werkzeug.utils import secure_filename
|
||||
from common.models.document import Document, DocumentVersion
|
||||
from common.models.document import Document, DocumentVersion, Catalog
|
||||
from common.extensions import db, minio_client
|
||||
from common.utils.celery_utils import current_celery
|
||||
from flask import current_app
|
||||
from flask_security import current_user
|
||||
import requests
|
||||
from urllib.parse import urlparse, unquote
|
||||
from urllib.parse import urlparse, unquote, urlunparse, parse_qs
|
||||
import os
|
||||
from .eveai_exceptions import EveAIInvalidLanguageException, EveAIDoubleURLException, EveAIUnsupportedFileType
|
||||
|
||||
from .config_field_types import normalize_json_field
|
||||
from .eveai_exceptions import (EveAIInvalidLanguageException, EveAIDoubleURLException, EveAIUnsupportedFileType,
|
||||
EveAIInvalidCatalog, EveAIInvalidDocument, EveAIInvalidDocumentVersion, EveAIException)
|
||||
from ..models.user import Tenant
|
||||
from common.utils.model_logging_utils import set_logging_information, update_logging_information
|
||||
from common.services.entitlements import LicenseUsageServices
|
||||
|
||||
MB_CONVERTOR = 1_048_576
|
||||
|
||||
|
||||
def get_file_size(file):
|
||||
try:
|
||||
# Als file een bytes object is of iets anders dat len() ondersteunt
|
||||
file_size = len(file)
|
||||
except TypeError:
|
||||
# Als file een FileStorage object is
|
||||
current_position = file.tell()
|
||||
file.seek(0, os.SEEK_END)
|
||||
file_size = file.tell()
|
||||
file.seek(current_position)
|
||||
|
||||
return file_size
|
||||
|
||||
|
||||
def create_document_stack(api_input, file, filename, extension, tenant_id):
|
||||
# Precheck if we can add a document to the stack
|
||||
|
||||
LicenseUsageServices.check_storage_and_embedding_quota(tenant_id, get_file_size(file)/MB_CONVERTOR)
|
||||
|
||||
# Create the Document
|
||||
catalog_id = int(api_input.get('catalog_id'))
|
||||
catalog = Catalog.query.get(catalog_id)
|
||||
if not catalog:
|
||||
raise EveAIInvalidCatalog(tenant_id, catalog_id)
|
||||
new_doc = create_document(api_input, filename, catalog_id)
|
||||
db.session.add(new_doc)
|
||||
|
||||
url = api_input.get('url', '')
|
||||
if url != '':
|
||||
url = cope_with_local_url(api_input.get('url', ''))
|
||||
|
||||
# Create the DocumentVersion
|
||||
new_doc_vers = create_version_for_document(new_doc, tenant_id,
|
||||
api_input.get('url', ''),
|
||||
url,
|
||||
api_input.get('sub_file_type', ''),
|
||||
api_input.get('language', 'en'),
|
||||
api_input.get('user_context', ''),
|
||||
api_input.get('user_metadata'),
|
||||
@@ -64,7 +96,8 @@ def create_document(form, filename, catalog_id):
|
||||
return new_doc
|
||||
|
||||
|
||||
def create_version_for_document(document, tenant_id, url, language, user_context, user_metadata, catalog_properties):
|
||||
def create_version_for_document(document, tenant_id, url, sub_file_type, language, user_context, user_metadata,
|
||||
catalog_properties):
|
||||
new_doc_vers = DocumentVersion()
|
||||
if url != '':
|
||||
new_doc_vers.url = url
|
||||
@@ -78,17 +111,18 @@ def create_version_for_document(document, tenant_id, url, language, user_context
|
||||
new_doc_vers.user_context = user_context
|
||||
|
||||
if user_metadata != '' and user_metadata is not None:
|
||||
new_doc_vers.user_metadata = user_metadata
|
||||
new_doc_vers.user_metadata = normalize_json_field(user_metadata, "user_metadata")
|
||||
|
||||
if catalog_properties != '' and catalog_properties is not None:
|
||||
new_doc_vers.catalog_properties = catalog_properties
|
||||
new_doc_vers.catalog_properties = normalize_json_field(catalog_properties, "catalog_properties")
|
||||
|
||||
if sub_file_type != '':
|
||||
new_doc_vers.sub_file_type = sub_file_type
|
||||
|
||||
new_doc_vers.document = document
|
||||
|
||||
set_logging_information(new_doc_vers, dt.now(tz.utc))
|
||||
|
||||
mark_tenant_storage_dirty(tenant_id)
|
||||
|
||||
return new_doc_vers
|
||||
|
||||
|
||||
@@ -109,7 +143,7 @@ def upload_file_for_version(doc_vers, file, extension, tenant_id):
|
||||
)
|
||||
doc_vers.bucket_name = bn
|
||||
doc_vers.object_name = on
|
||||
doc_vers.file_size = size / 1048576 # Convert bytes to MB
|
||||
doc_vers.file_size = size / MB_CONVERTOR # Convert bytes to MB
|
||||
|
||||
db.session.commit()
|
||||
current_app.logger.info(f'Successfully saved document to MinIO for tenant {tenant_id} for '
|
||||
@@ -121,35 +155,6 @@ def upload_file_for_version(doc_vers, file, extension, tenant_id):
|
||||
raise
|
||||
|
||||
|
||||
def set_logging_information(obj, timestamp):
|
||||
obj.created_at = timestamp
|
||||
obj.updated_at = timestamp
|
||||
|
||||
user_id = get_current_user_id()
|
||||
if user_id:
|
||||
obj.created_by = user_id
|
||||
obj.updated_by = user_id
|
||||
|
||||
|
||||
def update_logging_information(obj, timestamp):
|
||||
obj.updated_at = timestamp
|
||||
|
||||
user_id = get_current_user_id()
|
||||
if user_id:
|
||||
obj.updated_by = user_id
|
||||
|
||||
|
||||
def get_current_user_id():
|
||||
try:
|
||||
if current_user and current_user.is_authenticated:
|
||||
return current_user.id
|
||||
else:
|
||||
return None
|
||||
except Exception:
|
||||
# This will catch any errors if current_user is not available (e.g., in API context)
|
||||
return None
|
||||
|
||||
|
||||
def get_extension_from_content_type(content_type):
|
||||
content_type_map = {
|
||||
'text/html': 'html',
|
||||
@@ -163,6 +168,8 @@ def get_extension_from_content_type(content_type):
|
||||
|
||||
|
||||
def process_url(url, tenant_id):
|
||||
url = cope_with_local_url(url)
|
||||
|
||||
response = requests.head(url, allow_redirects=True)
|
||||
content_type = response.headers.get('Content-Type', '').split(';')[0]
|
||||
|
||||
@@ -194,36 +201,22 @@ def process_url(url, tenant_id):
|
||||
return file_content, filename, extension
|
||||
|
||||
|
||||
def process_multiple_urls(urls, tenant_id, api_input):
|
||||
results = []
|
||||
for url in urls:
|
||||
try:
|
||||
file_content, filename, extension = process_url(url, tenant_id)
|
||||
def clean_url(url):
|
||||
tracking_params = {"utm_source", "utm_medium", "utm_campaign", "utm_term", "utm_content",
|
||||
"hsa_acc", "hsa_cam", "hsa_grp", "hsa_ad", "hsa_src", "hsa_tgt", "hsa_kw",
|
||||
"hsa_mt", "hsa_net", "hsa_ver", "gad_source", "gbraid"}
|
||||
|
||||
url_input = api_input.copy()
|
||||
url_input.update({
|
||||
'url': url,
|
||||
'name': f"{api_input['name']}-{filename}" if api_input['name'] else filename
|
||||
})
|
||||
parsed_url = urlparse(url)
|
||||
query_params = parse_qs(parsed_url.query)
|
||||
|
||||
new_doc, new_doc_vers = create_document_stack(url_input, file_content, filename, extension, tenant_id)
|
||||
task_id = start_embedding_task(tenant_id, new_doc_vers.id)
|
||||
# Remove tracking params
|
||||
clean_params = {k: v for k, v in query_params.items() if k not in tracking_params}
|
||||
|
||||
results.append({
|
||||
'url': url,
|
||||
'document_id': new_doc.id,
|
||||
'document_version_id': new_doc_vers.id,
|
||||
'task_id': task_id,
|
||||
'status': 'success'
|
||||
})
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error processing URL {url}: {str(e)}")
|
||||
results.append({
|
||||
'url': url,
|
||||
'status': 'error',
|
||||
'message': str(e)
|
||||
})
|
||||
return results
|
||||
# Reconstruct the URL
|
||||
clean_query = "&".join(f"{k}={v[0]}" for k, v in clean_params.items()) if clean_params else ""
|
||||
cleaned_url = urlunparse(parsed_url._replace(query=clean_query))
|
||||
|
||||
return cleaned_url
|
||||
|
||||
|
||||
def start_embedding_task(tenant_id, doc_vers_id):
|
||||
@@ -236,14 +229,6 @@ def start_embedding_task(tenant_id, doc_vers_id):
|
||||
return task.id
|
||||
|
||||
|
||||
def validate_file_type(extension):
|
||||
current_app.logger.debug(f'Validating file type {extension}')
|
||||
current_app.logger.debug(f'Supported file types: {current_app.config["SUPPORTED_FILE_TYPES"]}')
|
||||
if extension not in current_app.config['SUPPORTED_FILE_TYPES']:
|
||||
raise EveAIUnsupportedFileType(f"Filetype {extension} is currently not supported. "
|
||||
f"Supported filetypes: {', '.join(current_app.config['SUPPORTED_FILE_TYPES'])}")
|
||||
|
||||
|
||||
def get_filename_from_url(url):
|
||||
parsed_url = urlparse(url)
|
||||
path_parts = parsed_url.path.split('/')
|
||||
@@ -261,11 +246,16 @@ def get_documents_list(page, per_page):
|
||||
return pagination
|
||||
|
||||
|
||||
def edit_document(document_id, name, valid_from, valid_to):
|
||||
doc = Document.query.get_or_404(document_id)
|
||||
doc.name = name
|
||||
doc.valid_from = valid_from
|
||||
doc.valid_to = valid_to
|
||||
def edit_document(tenant_id, document_id, name, valid_from, valid_to):
|
||||
doc = Document.query.get(document_id)
|
||||
if not doc:
|
||||
raise EveAIInvalidDocument(tenant_id, document_id)
|
||||
if name:
|
||||
doc.name = name
|
||||
if valid_from:
|
||||
doc.valid_from = valid_from
|
||||
if valid_to:
|
||||
doc.valid_to = valid_to
|
||||
update_logging_information(doc, dt.now(tz.utc))
|
||||
|
||||
try:
|
||||
@@ -277,10 +267,13 @@ def edit_document(document_id, name, valid_from, valid_to):
|
||||
return None, str(e)
|
||||
|
||||
|
||||
def edit_document_version(version_id, user_context, catalog_properties):
|
||||
doc_vers = DocumentVersion.query.get_or_404(version_id)
|
||||
def edit_document_version(tenant_id, version_id, user_context, catalog_properties):
|
||||
doc_vers = DocumentVersion.query.get(version_id)
|
||||
if not doc_vers:
|
||||
raise EveAIInvalidDocumentVersion(tenant_id, version_id)
|
||||
doc_vers.user_context = user_context
|
||||
doc_vers.catalog_properties = catalog_properties
|
||||
doc_vers.catalog_properties = normalize_json_field(catalog_properties, "catalog_properties")
|
||||
|
||||
update_logging_information(doc_vers, dt.now(tz.utc))
|
||||
|
||||
try:
|
||||
@@ -293,15 +286,20 @@ def edit_document_version(version_id, user_context, catalog_properties):
|
||||
|
||||
|
||||
def refresh_document_with_info(doc_id, tenant_id, api_input):
|
||||
doc = Document.query.get_or_404(doc_id)
|
||||
doc = Document.query.get(doc_id)
|
||||
if not doc:
|
||||
raise EveAIInvalidDocument(tenant_id, doc_id)
|
||||
old_doc_vers = DocumentVersion.query.filter_by(doc_id=doc_id).order_by(desc(DocumentVersion.id)).first()
|
||||
|
||||
if not old_doc_vers.url:
|
||||
return None, "This document has no URL. Only documents with a URL can be refreshed."
|
||||
|
||||
# Precheck if we have enough quota for the new version
|
||||
LicenseUsageServices.check_storage_and_embedding_quota(tenant_id, old_doc_vers.file_size)
|
||||
|
||||
new_doc_vers = create_version_for_document(
|
||||
doc, tenant_id,
|
||||
old_doc_vers.url,
|
||||
old_doc_vers.sub_file_type,
|
||||
api_input.get('language', old_doc_vers.language),
|
||||
api_input.get('user_context', old_doc_vers.user_context),
|
||||
api_input.get('user_metadata', old_doc_vers.user_metadata),
|
||||
@@ -317,11 +315,12 @@ def refresh_document_with_info(doc_id, tenant_id, api_input):
|
||||
db.session.rollback()
|
||||
return None, str(e)
|
||||
|
||||
response = requests.head(old_doc_vers.url, allow_redirects=True)
|
||||
url = cope_with_local_url(old_doc_vers.url)
|
||||
response = requests.head(url, allow_redirects=True)
|
||||
content_type = response.headers.get('Content-Type', '').split(';')[0]
|
||||
extension = get_extension_from_content_type(content_type)
|
||||
|
||||
response = requests.get(old_doc_vers.url)
|
||||
response = requests.get(url)
|
||||
response.raise_for_status()
|
||||
file_content = response.content
|
||||
|
||||
@@ -334,6 +333,59 @@ def refresh_document_with_info(doc_id, tenant_id, api_input):
|
||||
return new_doc_vers, task.id
|
||||
|
||||
|
||||
def refresh_document_with_content(doc_id: int, tenant_id: int, file_content: bytes, api_input: dict) -> tuple:
|
||||
"""
|
||||
Refresh document with new content
|
||||
|
||||
Args:
|
||||
doc_id: Document ID
|
||||
tenant_id: Tenant ID
|
||||
file_content: New file content
|
||||
api_input: Additional document information
|
||||
|
||||
Returns:
|
||||
Tuple of (new_version, task_id)
|
||||
"""
|
||||
doc = Document.query.get(doc_id)
|
||||
if not doc:
|
||||
raise EveAIInvalidDocument(tenant_id, doc_id)
|
||||
|
||||
old_doc_vers = DocumentVersion.query.filter_by(doc_id=doc_id).order_by(desc(DocumentVersion.id)).first()
|
||||
|
||||
# Precheck if we have enough quota for the new version
|
||||
LicenseUsageServices.check_storage_and_embedding_quota(tenant_id, get_file_size(file_content) / MB_CONVERTOR)
|
||||
|
||||
# Create new version with same file type as original
|
||||
extension = old_doc_vers.file_type
|
||||
|
||||
new_doc_vers = create_version_for_document(
|
||||
doc, tenant_id,
|
||||
'', # No URL for content-based updates
|
||||
old_doc_vers.sub_file_type,
|
||||
api_input.get('language', old_doc_vers.language),
|
||||
api_input.get('user_context', old_doc_vers.user_context),
|
||||
api_input.get('user_metadata', old_doc_vers.user_metadata),
|
||||
api_input.get('catalog_properties', old_doc_vers.catalog_properties),
|
||||
)
|
||||
|
||||
try:
|
||||
db.session.add(new_doc_vers)
|
||||
db.session.commit()
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
return None, str(e)
|
||||
|
||||
# Upload new content
|
||||
upload_file_for_version(new_doc_vers, file_content, extension, tenant_id)
|
||||
|
||||
# Start embedding task
|
||||
task = current_celery.send_task('create_embeddings', args=[tenant_id, new_doc_vers.id], queue='embeddings')
|
||||
current_app.logger.info(f'Embedding creation started for document {doc_id} on version {new_doc_vers.id} '
|
||||
f'with task id: {task.id}.')
|
||||
|
||||
return new_doc_vers, task.id
|
||||
|
||||
|
||||
# Update the existing refresh_document function to use the new refresh_document_with_info
|
||||
def refresh_document(doc_id, tenant_id):
|
||||
current_app.logger.info(f'Refreshing document {doc_id}')
|
||||
@@ -350,10 +402,70 @@ def refresh_document(doc_id, tenant_id):
|
||||
return refresh_document_with_info(doc_id, tenant_id, api_input)
|
||||
|
||||
|
||||
# Function triggered when a document_version is created or updated
|
||||
def mark_tenant_storage_dirty(tenant_id):
|
||||
tenant = db.session.query(Tenant).filter_by(id=int(tenant_id)).first()
|
||||
tenant.storage_dirty = True
|
||||
db.session.commit()
|
||||
def cope_with_local_url(url):
|
||||
parsed_url = urlparse(url)
|
||||
# Check if this is an internal WordPress URL (TESTING) and rewrite it
|
||||
if parsed_url.netloc in [current_app.config['EXTERNAL_WORDPRESS_BASE_URL']]:
|
||||
parsed_url = parsed_url._replace(
|
||||
scheme=current_app.config['WORDPRESS_PROTOCOL'],
|
||||
netloc=f"{current_app.config['WORDPRESS_HOST']}:{current_app.config['WORDPRESS_PORT']}"
|
||||
)
|
||||
url = urlunparse(parsed_url)
|
||||
|
||||
return url
|
||||
|
||||
|
||||
def lookup_document(tenant_id: int, lookup_criteria: dict, metadata_type: str) -> tuple[Document, DocumentVersion]:
|
||||
"""
|
||||
Look up a document using metadata criteria
|
||||
|
||||
Args:
|
||||
tenant_id: ID of the tenant
|
||||
lookup_criteria: Dictionary of key-value pairs to match in metadata
|
||||
metadata_type: Which metadata to search in ('user_metadata' or 'system_metadata')
|
||||
|
||||
Returns:
|
||||
Tuple of (Document, DocumentVersion) if found
|
||||
|
||||
Raises:
|
||||
ValueError: If invalid metadata_type provided
|
||||
EveAIException: If lookup fails
|
||||
"""
|
||||
if metadata_type not in ['user_metadata', 'system_metadata']:
|
||||
raise ValueError(f"Invalid metadata_type: {metadata_type}")
|
||||
|
||||
try:
|
||||
# Query for the latest document version matching the criteria
|
||||
query = (db.session.query(Document, DocumentVersion)
|
||||
.join(DocumentVersion)
|
||||
.filter(Document.id == DocumentVersion.doc_id)
|
||||
.order_by(DocumentVersion.id.desc()))
|
||||
|
||||
# Add metadata filtering using PostgreSQL JSONB operators
|
||||
metadata_field = getattr(DocumentVersion, metadata_type)
|
||||
for key, value in lookup_criteria.items():
|
||||
query = query.filter(metadata_field[key].astext == str(value))
|
||||
|
||||
# Get first result
|
||||
result = query.first()
|
||||
|
||||
if not result:
|
||||
raise EveAIException(
|
||||
f"No document found matching criteria in {metadata_type}",
|
||||
status_code=404
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
current_app.logger.error(f'Database error during document lookup for tenant {tenant_id}: {e}')
|
||||
raise EveAIException(
|
||||
"Database error during document lookup",
|
||||
status_code=500
|
||||
)
|
||||
except Exception as e:
|
||||
current_app.logger.error(f'Error during document lookup for tenant {tenant_id}: {e}')
|
||||
raise EveAIException(
|
||||
"Error during document lookup",
|
||||
status_code=500
|
||||
)
|
||||
|
||||
44
common/utils/dynamic_field_utils.py
Normal file
44
common/utils/dynamic_field_utils.py
Normal file
@@ -0,0 +1,44 @@
|
||||
def create_default_config_from_type_config(type_config):
|
||||
"""
|
||||
Creëert een dictionary met standaardwaarden gebaseerd op een typedefinitie configuratie.
|
||||
|
||||
Args:
|
||||
type_config (dict): Het configuration-veld van een typedefinitie (bijv. uit processor_types).
|
||||
|
||||
Returns:
|
||||
dict: Een dictionary met de naam van ieder veld als sleutel en de standaardwaarde als waarde.
|
||||
Alleen velden met een standaardwaarde of die verplicht zijn, worden opgenomen.
|
||||
|
||||
Voorbeeld:
|
||||
>>> config = PROCESSOR_TYPES["HTML_PROCESSOR"]["configuration"]
|
||||
>>> create_default_config_from_type_def(config)
|
||||
{'html_tags': 'p, h1, h2, h3, h4, h5, h6, li, table, thead, tbody, tr, td',
|
||||
'html_end_tags': 'p, li, table',
|
||||
'html_excluded_classes': '',
|
||||
'html_excluded_elements': 'header, footer, nav, script',
|
||||
'html_included_elements': 'article, main',
|
||||
'chunking_heading_level': 2}
|
||||
"""
|
||||
if not type_config:
|
||||
return {}
|
||||
|
||||
default_config = {}
|
||||
|
||||
for field_name, field_def in type_config.items():
|
||||
# Als het veld een standaardwaarde heeft, voeg deze toe
|
||||
if "default" in field_def:
|
||||
default_config[field_name] = field_def["default"]
|
||||
# Als het veld verplicht is maar geen standaardwaarde heeft, voeg een lege string toe
|
||||
elif field_def.get("required", False):
|
||||
# Kies een geschikte "lege" waarde op basis van het type
|
||||
field_type = field_def.get("type", "string")
|
||||
if field_type == "string":
|
||||
default_config[field_name] = ""
|
||||
elif field_type == "integer":
|
||||
default_config[field_name] = 0
|
||||
elif field_type == "boolean":
|
||||
default_config[field_name] = False
|
||||
else:
|
||||
default_config[field_name] = ""
|
||||
|
||||
return default_config
|
||||
@@ -10,8 +10,12 @@ class EveAIException(Exception):
|
||||
def to_dict(self):
|
||||
rv = dict(self.payload or ())
|
||||
rv['message'] = self.message
|
||||
rv['error'] = self.__class__.__name__
|
||||
return rv
|
||||
|
||||
def __str__(self):
|
||||
return self.message # Return the message when the exception is converted to a string
|
||||
|
||||
|
||||
class EveAIInvalidLanguageException(EveAIException):
|
||||
"""Raised when an invalid language is provided"""
|
||||
@@ -41,3 +45,206 @@ class EveAINoLicenseForTenant(EveAIException):
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAITenantNotFound(EveAIException):
|
||||
"""Raised when a tenant is not found"""
|
||||
|
||||
def __init__(self, tenant_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
message = f"Tenant {tenant_id} not found"
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAITenantInvalid(EveAIException):
|
||||
"""Raised when a tenant is invalid"""
|
||||
|
||||
def __init__(self, tenant_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' is not valid. Please contact the System Administrator."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAINoActiveLicense(EveAIException):
|
||||
"""Raised when a tenant has no active licenses"""
|
||||
|
||||
def __init__(self, tenant_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' has no active licenses. Please contact the System Administrator."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIInvalidCatalog(EveAIException):
|
||||
"""Raised when a catalog cannot be found"""
|
||||
|
||||
def __init__(self, tenant_id, catalog_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
self.catalog_id = catalog_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' has no valid catalog with ID {catalog_id}. Please contact the System Administrator."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIInvalidProcessor(EveAIException):
|
||||
"""Raised when no valid processor can be found for a given Catalog ID"""
|
||||
|
||||
def __init__(self, tenant_id, catalog_id, file_type, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
self.catalog_id = catalog_id
|
||||
self.file_type = file_type
|
||||
# Construct the message dynamically
|
||||
message = (f"Tenant with ID '{tenant_id}' has no valid {file_type} processor for catalog with ID {catalog_id}. "
|
||||
f"Please contact the System Administrator.")
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIInvalidDocument(EveAIException):
|
||||
"""Raised when a tenant has no document with given ID"""
|
||||
|
||||
def __init__(self, tenant_id, document_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
self.document_id = document_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' has no document with ID {document_id}."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIInvalidDocumentVersion(EveAIException):
|
||||
"""Raised when a tenant has no document version with given ID"""
|
||||
|
||||
def __init__(self, tenant_id, document_version_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
self.document_version_id = document_version_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' has no document version with ID {document_version_id}."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAISocketInputException(EveAIException):
|
||||
"""Raised when a socket call receives an invalid payload"""
|
||||
|
||||
def __init__(self, message, status_code=400, payload=None):
|
||||
super.__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIInvalidEmbeddingModel(EveAIException):
|
||||
"""Raised when no or an invalid embedding model is provided in the catalog"""
|
||||
|
||||
def __init__(self, tenant_id, catalog_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
self.catalog_id = catalog_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' has no or an invalid embedding model in Catalog {catalog_id}."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIDoublePartner(EveAIException):
|
||||
"""Raised when there is already a partner defined for a given tenant (while registering a partner)"""
|
||||
|
||||
def __init__(self, tenant_id, status_code=400, payload=None):
|
||||
self.tenant_id = tenant_id
|
||||
# Construct the message dynamically
|
||||
message = f"Tenant with ID '{tenant_id}' is already defined as a Partner."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIRoleAssignmentException(EveAIException):
|
||||
"""Exception raised when a role cannot be assigned due to business rules"""
|
||||
|
||||
def __init__(self, message, status_code=403, payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAINoManagementPartnerService(EveAIException):
|
||||
"""Exception raised when the operation requires the logged in partner (or selected parter by Super User)
|
||||
does not have a MANAGEMENT_SERVICE"""
|
||||
|
||||
def __init__(self, message="No Management Service defined for partner", status_code=403, payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAINoSessionTenant(EveAIException):
|
||||
"""Exception raised when no session tenant is set"""
|
||||
|
||||
def __init__(self, message="No Session Tenant selected. Cannot perform requested action.", status_code=403,
|
||||
payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAINoSessionPartner(EveAIException):
|
||||
"""Exception raised when no session partner is set"""
|
||||
|
||||
def __init__(self, message="No Session Partner selected. Cannot perform requested action.", status_code=403,
|
||||
payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAINoManagementPartnerForTenant(EveAIException):
|
||||
"""Exception raised when the selected partner is no management partner for tenant"""
|
||||
|
||||
def __init__(self, message="No Management Partner for Tenant", status_code=403, payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIQuotaExceeded(EveAIException):
|
||||
"""Base exception for quota-related errors"""
|
||||
|
||||
def __init__(self, message, quota_type, current_usage, limit, additional=0, status_code=400, payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
self.quota_type = quota_type
|
||||
self.current_usage = current_usage
|
||||
self.limit = limit
|
||||
self.additional = additional
|
||||
|
||||
|
||||
class EveAIStorageQuotaExceeded(EveAIQuotaExceeded):
|
||||
"""Raised when storage quota is exceeded"""
|
||||
|
||||
def __init__(self, current_usage, limit, additional, status_code=400, payload=None):
|
||||
message = (f"Storage quota exceeded. Current: {current_usage:.1f}MB, "
|
||||
f"Additional: {additional:.1f}MB, Limit: {limit}MB")
|
||||
super().__init__(message, "storage", current_usage, limit, additional, status_code, payload)
|
||||
|
||||
|
||||
class EveAIEmbeddingQuotaExceeded(EveAIQuotaExceeded):
|
||||
"""Raised when embedding quota is exceeded"""
|
||||
|
||||
def __init__(self, current_usage, limit, additional, status_code=400, payload=None):
|
||||
message = (f"Embedding quota exceeded. Current: {current_usage:.1f}MB, "
|
||||
f"Additional: {additional:.1f}MB, Limit: {limit}MB")
|
||||
super().__init__(message, "embedding", current_usage, limit, additional, status_code, payload)
|
||||
|
||||
|
||||
class EveAIInteractionQuotaExceeded(EveAIQuotaExceeded):
|
||||
"""Raised when the interaction token quota is exceeded"""
|
||||
|
||||
def __init__(self, current_usage, limit, status_code=400, payload=None):
|
||||
message = (f"Interaction token quota exceeded. Current: {current_usage:.2f}M tokens, "
|
||||
f"Limit: {limit:.2f}M tokens")
|
||||
super().__init__(message, "interaction", current_usage, limit, 0, status_code, payload)
|
||||
|
||||
|
||||
class EveAIQuotaWarning(EveAIException):
|
||||
"""Warning when approaching quota limits (not blocking)"""
|
||||
|
||||
def __init__(self, message, quota_type, usage_percentage, status_code=200, payload=None):
|
||||
super().__init__(message, status_code, payload)
|
||||
self.quota_type = quota_type
|
||||
self.usage_percentage = usage_percentage
|
||||
|
||||
|
||||
class EveAILicensePeriodsExceeded(EveAIException):
|
||||
"""Raised when no more license periods can be created for a given license"""
|
||||
|
||||
def __init__(self, license_id, status_code=400, payload=None):
|
||||
message = f"No more license periods can be created for license with ID {license_id}. "
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
class EveAIPendingLicensePeriod(EveAIException):
|
||||
"""Raised when a license period is pending"""
|
||||
|
||||
def __init__(self, status_code=400, payload=None):
|
||||
message = f"Basic Fee Payment has not been received yet. Please ensure payment has been made, and please wait for payment to be processed."
|
||||
super().__init__(message, status_code, payload)
|
||||
|
||||
|
||||
112
common/utils/execution_progress.py
Normal file
112
common/utils/execution_progress.py
Normal file
@@ -0,0 +1,112 @@
|
||||
# common/utils/execution_progress.py
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
from typing import Generator
|
||||
from redis import Redis, RedisError
|
||||
import json
|
||||
from flask import current_app
|
||||
|
||||
|
||||
class ExecutionProgressTracker:
|
||||
"""Tracks progress of specialist executions using Redis"""
|
||||
|
||||
def __init__(self):
|
||||
try:
|
||||
redis_url = current_app.config['SPECIALIST_EXEC_PUBSUB']
|
||||
|
||||
self.redis = Redis.from_url(redis_url, socket_timeout=5)
|
||||
# Test the connection
|
||||
self.redis.ping()
|
||||
|
||||
self.expiry = 3600 # 1 hour expiry
|
||||
except RedisError as e:
|
||||
current_app.logger.error(f"Failed to connect to Redis: {str(e)}")
|
||||
raise
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Unexpected error during Redis initialization: {str(e)}")
|
||||
raise
|
||||
|
||||
def _get_key(self, execution_id: str) -> str:
|
||||
return f"specialist_execution:{execution_id}"
|
||||
|
||||
def send_update(self, ctask_id: str, processing_type: str, data: dict):
|
||||
"""Send an update about execution progress"""
|
||||
try:
|
||||
key = self._get_key(ctask_id)
|
||||
|
||||
# First verify Redis is still connected
|
||||
try:
|
||||
self.redis.ping()
|
||||
except RedisError:
|
||||
current_app.logger.error("Lost Redis connection. Attempting to reconnect...")
|
||||
self.__init__() # Reinitialize connection
|
||||
|
||||
update = {
|
||||
'processing_type': processing_type,
|
||||
'data': data,
|
||||
'timestamp': dt.now(tz=tz.utc)
|
||||
}
|
||||
|
||||
# Log initial state
|
||||
try:
|
||||
orig_len = self.redis.llen(key)
|
||||
|
||||
# Try to serialize the update and check the result
|
||||
try:
|
||||
serialized_update = json.dumps(update, default=str) # Add default handler for datetime
|
||||
except TypeError as e:
|
||||
current_app.logger.error(f"Failed to serialize update: {str(e)}")
|
||||
raise
|
||||
|
||||
# Store update in list with pipeline for atomicity
|
||||
with self.redis.pipeline() as pipe:
|
||||
pipe.rpush(key, serialized_update)
|
||||
pipe.publish(key, serialized_update)
|
||||
pipe.expire(key, self.expiry)
|
||||
results = pipe.execute()
|
||||
|
||||
new_len = self.redis.llen(key)
|
||||
|
||||
if new_len <= orig_len:
|
||||
current_app.logger.error(
|
||||
f"List length did not increase as expected. Original: {orig_len}, New: {new_len}")
|
||||
|
||||
except RedisError as e:
|
||||
current_app.logger.error(f"Redis operation failed: {str(e)}")
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Unexpected error in send_update: {str(e)}, type: {type(e)}")
|
||||
raise
|
||||
|
||||
def get_updates(self, ctask_id: str) -> Generator[str, None, None]:
|
||||
key = self._get_key(ctask_id)
|
||||
pubsub = self.redis.pubsub()
|
||||
pubsub.subscribe(key)
|
||||
|
||||
try:
|
||||
# First yield any existing updates
|
||||
length = self.redis.llen(key)
|
||||
if length > 0:
|
||||
updates = self.redis.lrange(key, 0, -1)
|
||||
for update in updates:
|
||||
update_data = json.loads(update.decode('utf-8'))
|
||||
# Use processing_type for the event
|
||||
yield f"event: {update_data['processing_type']}\n"
|
||||
yield f"data: {json.dumps(update_data)}\n\n"
|
||||
|
||||
# Then listen for new updates
|
||||
while True:
|
||||
message = pubsub.get_message(timeout=30) # message['type'] is Redis pub/sub type
|
||||
if message is None:
|
||||
yield ": keepalive\n\n"
|
||||
continue
|
||||
|
||||
if message['type'] == 'message': # This is Redis pub/sub type
|
||||
update_data = json.loads(message['data'].decode('utf-8'))
|
||||
yield f"data: {message['data'].decode('utf-8')}\n\n"
|
||||
|
||||
# Check processing_type for completion
|
||||
if update_data['processing_type'] in ['Task Complete', 'Task Error']:
|
||||
break
|
||||
finally:
|
||||
pubsub.unsubscribe()
|
||||
11
common/utils/form_assistants.py
Normal file
11
common/utils/form_assistants.py
Normal file
@@ -0,0 +1,11 @@
|
||||
import json
|
||||
|
||||
from wtforms.validators import ValidationError
|
||||
|
||||
|
||||
def validate_json(form, field):
|
||||
if field.data:
|
||||
try:
|
||||
json.loads(field.data)
|
||||
except json.JSONDecodeError:
|
||||
raise ValidationError('Invalid JSON format')
|
||||
57
common/utils/log_utils.py
Normal file
57
common/utils/log_utils.py
Normal file
@@ -0,0 +1,57 @@
|
||||
import pandas as pd
|
||||
from sqlalchemy import inspect
|
||||
from typing import Any, List, Union, Optional
|
||||
|
||||
|
||||
def format_query_results(query_results: Any) -> str:
|
||||
"""
|
||||
Format query results as a readable string using pandas
|
||||
|
||||
Args:
|
||||
query_results: SQLAlchemy query, query results, or model instance(s)
|
||||
|
||||
Returns:
|
||||
Formatted string representation of the query results
|
||||
"""
|
||||
try:
|
||||
# If it's a query object, execute it
|
||||
if hasattr(query_results, 'all'):
|
||||
results = query_results.all()
|
||||
elif not isinstance(query_results, list):
|
||||
results = [query_results]
|
||||
else:
|
||||
results = query_results
|
||||
|
||||
# Handle different types of results
|
||||
if results and hasattr(results[0], '__table__'):
|
||||
# SQLAlchemy ORM objects
|
||||
data = []
|
||||
for item in results:
|
||||
row = {}
|
||||
for column in inspect(item).mapper.column_attrs:
|
||||
row[column.key] = getattr(item, column.key)
|
||||
data.append(row)
|
||||
df = pd.DataFrame(data)
|
||||
elif results and isinstance(results[0], tuple):
|
||||
# Join query results (tuples)
|
||||
if hasattr(results[0], '_fields'): # Named tuples
|
||||
df = pd.DataFrame(results)
|
||||
else:
|
||||
# Regular tuples - try to get column names from query
|
||||
if hasattr(query_results, 'statement'):
|
||||
columns = query_results.statement.columns.keys()
|
||||
df = pd.DataFrame(results, columns=columns)
|
||||
else:
|
||||
df = pd.DataFrame(results)
|
||||
else:
|
||||
# Fallback for other types
|
||||
df = pd.DataFrame(results)
|
||||
|
||||
# Format the output with pandas
|
||||
with pd.option_context('display.max_rows', 20, 'display.max_columns', None,
|
||||
'display.width', 1000):
|
||||
formatted_output = f"Query returned {len(df)} results:\n{df}"
|
||||
|
||||
return formatted_output
|
||||
except Exception as e:
|
||||
return f"Error formatting query results: {str(e)}"
|
||||
46
common/utils/mail_utils.py
Normal file
46
common/utils/mail_utils.py
Normal file
@@ -0,0 +1,46 @@
|
||||
from scaleway import Client
|
||||
from scaleway.tem.v1alpha1.api import TemV1Alpha1API
|
||||
from scaleway.tem.v1alpha1.types import CreateEmailRequestAddress
|
||||
from html2text import HTML2Text
|
||||
from flask import current_app
|
||||
|
||||
|
||||
def send_email(to_email, to_name, subject, html):
|
||||
current_app.logger.debug(f"Sending email to {to_email} with subject {subject}")
|
||||
access_key = current_app.config['SW_EMAIL_ACCESS_KEY']
|
||||
secret_key = current_app.config['SW_EMAIL_SECRET_KEY']
|
||||
default_project_id = current_app.config['SW_PROJECT']
|
||||
default_region = "fr-par"
|
||||
current_app.logger.debug(f"Access Key: {access_key}\nSecret Key: {secret_key}\n"
|
||||
f"Default Project ID: {default_project_id}\nDefault Region: {default_region}")
|
||||
client = Client(
|
||||
access_key=access_key,
|
||||
secret_key=secret_key,
|
||||
default_project_id=default_project_id,
|
||||
default_region=default_region
|
||||
)
|
||||
current_app.logger.debug(f"Scaleway Client Initialized")
|
||||
tem = TemV1Alpha1API(client)
|
||||
current_app.logger.debug(f"Tem Initialized")
|
||||
from_ = CreateEmailRequestAddress(email=current_app.config['SW_EMAIL_SENDER'],
|
||||
name=current_app.config['SW_EMAIL_NAME'])
|
||||
to_ = CreateEmailRequestAddress(email=to_email, name=to_name)
|
||||
|
||||
email = tem.create_email(
|
||||
from_=from_,
|
||||
to=[to_],
|
||||
subject=subject,
|
||||
text=html_to_text(html),
|
||||
html=html,
|
||||
project_id=default_project_id,
|
||||
)
|
||||
current_app.logger.debug(f"Email sent to {to_email}")
|
||||
|
||||
|
||||
def html_to_text(html_content):
|
||||
"""Convert HTML to plain text using html2text"""
|
||||
h = HTML2Text()
|
||||
h.ignore_images = True
|
||||
h.ignore_emphasis = False
|
||||
h.body_width = 0 # No wrapping
|
||||
return h.handle(html_content)
|
||||
@@ -4,10 +4,11 @@ for handling tenant requests
|
||||
"""
|
||||
|
||||
from flask_security import current_user
|
||||
from flask import session, current_app, redirect
|
||||
from common.utils.nginx_utils import prefixed_url_for
|
||||
|
||||
from flask import session
|
||||
from .database import Database
|
||||
from .eveai_exceptions import EveAINoSessionTenant, EveAINoSessionPartner, EveAINoManagementPartnerService, \
|
||||
EveAINoManagementPartnerForTenant
|
||||
from common.services.user import UserServices
|
||||
|
||||
|
||||
def mw_before_request():
|
||||
@@ -17,20 +18,27 @@ def mw_before_request():
|
||||
"""
|
||||
|
||||
if 'tenant' not in session:
|
||||
current_app.logger.warning('No tenant defined in session')
|
||||
return redirect(prefixed_url_for('security_bp.login'))
|
||||
raise EveAINoSessionTenant()
|
||||
|
||||
tenant_id = session['tenant']['id']
|
||||
if not tenant_id:
|
||||
raise Exception('Cannot switch schema for tenant: no tenant defined in session')
|
||||
raise EveAINoSessionTenant()
|
||||
|
||||
for role in current_user.roles:
|
||||
current_app.logger.debug(f'In middleware: User {current_user.email} has role {role.name}')
|
||||
switch_allowed = False
|
||||
if current_user.has_role('Super User'):
|
||||
switch_allowed = True
|
||||
if current_user.has_role('Tenant Admin') and current_user.tenant_id == tenant_id:
|
||||
switch_allowed = True
|
||||
if current_user.has_role('Partner Admin'):
|
||||
if 'partner' not in session:
|
||||
raise EveAINoSessionPartner()
|
||||
management_service = next((service for service in session['partner']['services']
|
||||
if service.get('type') == 'MANAGEMENT_SERVICE'), None)
|
||||
if not management_service:
|
||||
raise EveAINoManagementPartnerService()
|
||||
if not UserServices.can_user_edit_tenant(tenant_id):
|
||||
raise EveAINoManagementPartnerForTenant()
|
||||
|
||||
# user = User.query.get(current_user.id)
|
||||
if current_user.has_role('Super User') or current_user.tenant_id == tenant_id:
|
||||
Database(tenant_id).switch_schema()
|
||||
else:
|
||||
raise Exception(f'Cannot switch schema for tenant {tenant_id}: user {current_user.email} does not have access')
|
||||
Database(tenant_id).switch_schema()
|
||||
|
||||
|
||||
|
||||
@@ -33,6 +33,9 @@ class MinioClient:
|
||||
def generate_object_name(self, document_id, language, version_id, filename):
|
||||
return f"{document_id}/{language}/{version_id}/{filename}"
|
||||
|
||||
def generate_asset_name(self, asset_version_id, file_name, content_type):
|
||||
return f"assets/{asset_version_id}/{file_name}.{content_type}"
|
||||
|
||||
def upload_document_file(self, tenant_id, document_id, language, version_id, filename, file_data):
|
||||
bucket_name = self.generate_bucket_name(tenant_id)
|
||||
object_name = self.generate_object_name(document_id, language, version_id, filename)
|
||||
@@ -54,6 +57,26 @@ class MinioClient:
|
||||
except S3Error as err:
|
||||
raise Exception(f"Error occurred while uploading file: {err}")
|
||||
|
||||
def upload_asset_file(self, bucket_name, asset_version_id, file_name, file_type, file_data):
|
||||
object_name = self.generate_asset_name(asset_version_id, file_name, file_type)
|
||||
|
||||
try:
|
||||
if isinstance(file_data, FileStorage):
|
||||
file_data = file_data.read()
|
||||
elif isinstance(file_data, io.BytesIO):
|
||||
file_data = file_data.getvalue()
|
||||
elif isinstance(file_data, str):
|
||||
file_data = file_data.encode('utf-8')
|
||||
elif not isinstance(file_data, bytes):
|
||||
raise TypeError('Unsupported file type. Expected FileStorage, BytesIO, str, or bytes.')
|
||||
|
||||
self.client.put_object(
|
||||
bucket_name, object_name, io.BytesIO(file_data), len(file_data)
|
||||
)
|
||||
return object_name, len(file_data)
|
||||
except S3Error as err:
|
||||
raise Exception(f"Error occurred while uploading asset: {err}")
|
||||
|
||||
def download_document_file(self, tenant_id, bucket_name, object_name):
|
||||
try:
|
||||
response = self.client.get_object(bucket_name, object_name)
|
||||
|
||||
29
common/utils/model_logging_utils.py
Normal file
29
common/utils/model_logging_utils.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from flask_security import current_user
|
||||
|
||||
|
||||
def set_logging_information(obj, timestamp):
|
||||
obj.created_at = timestamp
|
||||
obj.updated_at = timestamp
|
||||
|
||||
user_id = get_current_user_id()
|
||||
if user_id:
|
||||
obj.created_by = user_id
|
||||
obj.updated_by = user_id
|
||||
|
||||
def update_logging_information(obj, timestamp):
|
||||
obj.updated_at = timestamp
|
||||
|
||||
user_id = get_current_user_id()
|
||||
if user_id:
|
||||
obj.updated_by = user_id
|
||||
|
||||
|
||||
def get_current_user_id():
|
||||
try:
|
||||
if current_user and current_user.is_authenticated:
|
||||
return current_user.id
|
||||
else:
|
||||
return None
|
||||
except Exception:
|
||||
# This will catch any errors if current_user is not available (e.g., in API context)
|
||||
return None
|
||||
@@ -1,249 +1,40 @@
|
||||
import os
|
||||
from typing import Dict, Any, Optional, Tuple
|
||||
|
||||
import langcodes
|
||||
from flask import current_app
|
||||
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_core.pydantic_v1 import BaseModel, Field
|
||||
from typing import List, Any, Iterator
|
||||
from collections.abc import MutableMapping
|
||||
from openai import OpenAI
|
||||
from portkey_ai import createHeaders, PORTKEY_GATEWAY_URL
|
||||
from portkey_ai.langchain.portkey_langchain_callback_handler import LangchainCallbackHandler
|
||||
from langchain_core.language_models import BaseChatModel
|
||||
|
||||
from common.langchain.llm_metrics_handler import LLMMetricsHandler
|
||||
from common.langchain.tracked_openai_embeddings import TrackedOpenAIEmbeddings
|
||||
from common.langchain.tracked_transcribe import tracked_transcribe
|
||||
from common.models.document import EmbeddingSmallOpenAI, EmbeddingLargeOpenAI, Catalog
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
from langchain_mistralai import ChatMistralAI
|
||||
from flask import current_app
|
||||
|
||||
from common.eveai_model.tracked_mistral_embeddings import TrackedMistralAIEmbeddings
|
||||
from common.langchain.tracked_transcription import TrackedOpenAITranscription
|
||||
from common.models.user import Tenant
|
||||
from config.model_config import MODEL_CONFIG
|
||||
from common.utils.business_event_context import current_event
|
||||
from common.extensions import cache_manager
|
||||
from common.models.document import EmbeddingMistral
|
||||
from common.utils.eveai_exceptions import EveAITenantNotFound, EveAIInvalidEmbeddingModel
|
||||
from crewai import LLM
|
||||
|
||||
embedding_llm_model_cache: Dict[Tuple[str, float], BaseChatModel] = {}
|
||||
crewai_llm_model_cache: Dict[Tuple[str, float], LLM] = {}
|
||||
llm_metrics_handler = LLMMetricsHandler()
|
||||
|
||||
|
||||
class CitedAnswer(BaseModel):
|
||||
"""Default docstring - to be replaced with actual prompt"""
|
||||
def create_language_template(template: str, language: str) -> str:
|
||||
"""
|
||||
Replace language placeholder in template with specified language
|
||||
|
||||
answer: str = Field(
|
||||
...,
|
||||
description="The answer to the user question, based on the given sources",
|
||||
)
|
||||
citations: List[int] = Field(
|
||||
...,
|
||||
description="The integer IDs of the SPECIFIC sources that were used to generate the answer"
|
||||
)
|
||||
insufficient_info: bool = Field(
|
||||
False, # Default value is set to False
|
||||
description="A boolean indicating wether given sources were sufficient or not to generate the answer"
|
||||
)
|
||||
Args:
|
||||
template: Template string with {language} placeholder
|
||||
language: Language code to insert
|
||||
|
||||
|
||||
def set_language_prompt_template(cls, language_prompt):
|
||||
cls.__doc__ = language_prompt
|
||||
|
||||
|
||||
class ModelVariables(MutableMapping):
|
||||
def __init__(self, tenant: Tenant, catalog_id=None):
|
||||
self.tenant = tenant
|
||||
self.catalog_id = catalog_id
|
||||
self._variables = self._initialize_variables()
|
||||
self._embedding_model = None
|
||||
self._llm = None
|
||||
self._llm_no_rag = None
|
||||
self._transcription_client = None
|
||||
self._prompt_templates = {}
|
||||
self._embedding_db_model = None
|
||||
self.llm_metrics_handler = LLMMetricsHandler()
|
||||
self._transcription_client = None
|
||||
|
||||
def _initialize_variables(self):
|
||||
variables = {}
|
||||
|
||||
# Get the Catalog if catalog_id is passed
|
||||
if self.catalog_id:
|
||||
catalog = Catalog.query.get_or_404(self.catalog_id)
|
||||
|
||||
# We initialize the variables that are available knowing the tenant.
|
||||
variables['embed_tuning'] = catalog.embed_tuning or False
|
||||
|
||||
# Set HTML Chunking Variables
|
||||
variables['html_tags'] = catalog.html_tags
|
||||
variables['html_end_tags'] = catalog.html_end_tags
|
||||
variables['html_included_elements'] = catalog.html_included_elements
|
||||
variables['html_excluded_elements'] = catalog.html_excluded_elements
|
||||
variables['html_excluded_classes'] = catalog.html_excluded_classes
|
||||
|
||||
# Set Chunk Size variables
|
||||
variables['min_chunk_size'] = catalog.min_chunk_size
|
||||
variables['max_chunk_size'] = catalog.max_chunk_size
|
||||
|
||||
# Set the RAG Context (will have to change once specialists are defined
|
||||
variables['rag_context'] = self.tenant.rag_context or " "
|
||||
# Temporary setting until we have Specialists
|
||||
variables['rag_tuning'] = False
|
||||
variables['RAG_temperature'] = 0.3
|
||||
variables['no_RAG_temperature'] = 0.5
|
||||
variables['k'] = 8
|
||||
variables['similarity_threshold'] = 0.4
|
||||
|
||||
# Set model providers
|
||||
variables['embedding_provider'], variables['embedding_model'] = self.tenant.embedding_model.rsplit('.', 1)
|
||||
variables['llm_provider'], variables['llm_model'] = self.tenant.llm_model.rsplit('.', 1)
|
||||
variables["templates"] = current_app.config['PROMPT_TEMPLATES'][(f"{variables['llm_provider']}."
|
||||
f"{variables['llm_model']}")]
|
||||
current_app.logger.info(f"Loaded prompt templates: \n")
|
||||
current_app.logger.info(f"{variables['templates']}")
|
||||
|
||||
# Set model-specific configurations
|
||||
model_config = MODEL_CONFIG.get(variables['llm_provider'], {}).get(variables['llm_model'], {})
|
||||
variables.update(model_config)
|
||||
|
||||
variables['annotation_chunk_length'] = current_app.config['ANNOTATION_TEXT_CHUNK_LENGTH'][self.tenant.llm_model]
|
||||
|
||||
if variables['tool_calling_supported']:
|
||||
variables['cited_answer_cls'] = CitedAnswer
|
||||
|
||||
variables['max_compression_duration'] = current_app.config['MAX_COMPRESSION_DURATION']
|
||||
variables['max_transcription_duration'] = current_app.config['MAX_TRANSCRIPTION_DURATION']
|
||||
variables['compression_cpu_limit'] = current_app.config['COMPRESSION_CPU_LIMIT']
|
||||
variables['compression_process_delay'] = current_app.config['COMPRESSION_PROCESS_DELAY']
|
||||
|
||||
return variables
|
||||
|
||||
@property
|
||||
def embedding_model(self):
|
||||
api_key = os.getenv('OPENAI_API_KEY')
|
||||
model = self._variables['embedding_model']
|
||||
self._embedding_model = TrackedOpenAIEmbeddings(api_key=api_key,
|
||||
model=model,
|
||||
)
|
||||
self._embedding_db_model = EmbeddingSmallOpenAI \
|
||||
if model == 'text-embedding-3-small' \
|
||||
else EmbeddingLargeOpenAI
|
||||
|
||||
return self._embedding_model
|
||||
|
||||
@property
|
||||
def llm(self):
|
||||
api_key = self.get_api_key_for_llm()
|
||||
self._llm = ChatOpenAI(api_key=api_key,
|
||||
model=self._variables['llm_model'],
|
||||
temperature=self._variables['RAG_temperature'],
|
||||
callbacks=[self.llm_metrics_handler])
|
||||
return self._llm
|
||||
|
||||
@property
|
||||
def llm_no_rag(self):
|
||||
api_key = self.get_api_key_for_llm()
|
||||
self._llm_no_rag = ChatOpenAI(api_key=api_key,
|
||||
model=self._variables['llm_model'],
|
||||
temperature=self._variables['RAG_temperature'],
|
||||
callbacks=[self.llm_metrics_handler])
|
||||
return self._llm_no_rag
|
||||
|
||||
def get_api_key_for_llm(self):
|
||||
if self._variables['llm_provider'] == 'openai':
|
||||
api_key = os.getenv('OPENAI_API_KEY')
|
||||
else: # self._variables['llm_provider'] == 'anthropic'
|
||||
api_key = os.getenv('ANTHROPIC_API_KEY')
|
||||
|
||||
return api_key
|
||||
|
||||
@property
|
||||
def transcription_client(self):
|
||||
api_key = os.getenv('OPENAI_API_KEY')
|
||||
self._transcription_client = OpenAI(api_key=api_key, )
|
||||
self._variables['transcription_model'] = 'whisper-1'
|
||||
return self._transcription_client
|
||||
|
||||
def transcribe(self, *args, **kwargs):
|
||||
return tracked_transcribe(self._transcription_client, *args, **kwargs)
|
||||
|
||||
@property
|
||||
def embedding_db_model(self):
|
||||
if self._embedding_db_model is None:
|
||||
self._embedding_db_model = self.get_embedding_db_model()
|
||||
return self._embedding_db_model
|
||||
|
||||
def get_embedding_db_model(self):
|
||||
current_app.logger.debug("In get_embedding_db_model")
|
||||
if self._embedding_db_model is None:
|
||||
self._embedding_db_model = EmbeddingSmallOpenAI \
|
||||
if self._variables['embedding_model'] == 'text-embedding-3-small' \
|
||||
else EmbeddingLargeOpenAI
|
||||
current_app.logger.debug(f"Embedding DB Model: {self._embedding_db_model}")
|
||||
return self._embedding_db_model
|
||||
|
||||
def get_prompt_template(self, template_name: str) -> str:
|
||||
current_app.logger.info(f"Getting prompt template for {template_name}")
|
||||
if template_name not in self._prompt_templates:
|
||||
self._prompt_templates[template_name] = self._load_prompt_template(template_name)
|
||||
return self._prompt_templates[template_name]
|
||||
|
||||
def _load_prompt_template(self, template_name: str) -> str:
|
||||
# In the future, this method will make an API call to Portkey
|
||||
# For now, we'll simulate it with a placeholder implementation
|
||||
# You can replace this with your current prompt loading logic
|
||||
return self._variables['templates'][template_name]
|
||||
|
||||
def __getitem__(self, key: str) -> Any:
|
||||
current_app.logger.debug(f"ModelVariables: Getting {key}")
|
||||
# Support older template names (suffix = _template)
|
||||
if key.endswith('_template'):
|
||||
key = key[:-len('_template')]
|
||||
current_app.logger.debug(f"ModelVariables: Getting modified {key}")
|
||||
if key == 'embedding_model':
|
||||
return self.embedding_model
|
||||
elif key == 'embedding_db_model':
|
||||
return self.embedding_db_model
|
||||
elif key == 'llm':
|
||||
return self.llm
|
||||
elif key == 'llm_no_rag':
|
||||
return self.llm_no_rag
|
||||
elif key == 'transcription_client':
|
||||
return self.transcription_client
|
||||
elif key in self._variables.get('prompt_templates', []):
|
||||
return self.get_prompt_template(key)
|
||||
else:
|
||||
value = self._variables.get(key)
|
||||
if value is not None:
|
||||
return value
|
||||
else:
|
||||
raise KeyError(f'Variable {key} does not exist in ModelVariables')
|
||||
|
||||
def __setitem__(self, key: str, value: Any) -> None:
|
||||
self._variables[key] = value
|
||||
|
||||
def __delitem__(self, key: str) -> None:
|
||||
del self._variables[key]
|
||||
|
||||
def __iter__(self) -> Iterator[str]:
|
||||
return iter(self._variables)
|
||||
|
||||
def __len__(self):
|
||||
return len(self._variables)
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
return self.__getitem__(key) or default
|
||||
|
||||
def update(self, **kwargs) -> None:
|
||||
self._variables.update(kwargs)
|
||||
|
||||
def items(self):
|
||||
return self._variables.items()
|
||||
|
||||
def keys(self):
|
||||
return self._variables.keys()
|
||||
|
||||
def values(self):
|
||||
return self._variables.values()
|
||||
|
||||
|
||||
def select_model_variables(tenant, catalog_id=None):
|
||||
model_variables = ModelVariables(tenant=tenant, catalog_id=catalog_id)
|
||||
return model_variables
|
||||
|
||||
|
||||
def create_language_template(template, language):
|
||||
Returns:
|
||||
str: Template with language placeholder replaced
|
||||
"""
|
||||
try:
|
||||
full_language = langcodes.Language.make(language=language)
|
||||
language_template = template.replace('{language}', full_language.display_name())
|
||||
@@ -253,5 +44,121 @@ def create_language_template(template, language):
|
||||
return language_template
|
||||
|
||||
|
||||
def replace_variable_in_template(template, variable, value):
|
||||
return template.replace(variable, value)
|
||||
def replace_variable_in_template(template: str, variable: str, value: str) -> str:
|
||||
"""
|
||||
Replace a variable placeholder in template with specified value
|
||||
|
||||
Args:
|
||||
template: Template string with variable placeholder
|
||||
variable: Variable placeholder to replace (e.g. "{tenant_context}")
|
||||
value: Value to insert
|
||||
|
||||
Returns:
|
||||
str: Template with variable placeholder replaced
|
||||
"""
|
||||
return template.replace(variable, value or "")
|
||||
|
||||
|
||||
def get_embedding_model_and_class(tenant_id, catalog_id, full_embedding_name="mistral.mistral-embed"):
|
||||
"""
|
||||
Retrieve the embedding model and embedding model class to store Embeddings
|
||||
|
||||
Args:
|
||||
tenant_id: ID of the tenant
|
||||
catalog_id: ID of the catalog
|
||||
full_embedding_name: The full name of the embedding model: <provider>.<model>
|
||||
|
||||
Returns:
|
||||
embedding_model, embedding_model_class
|
||||
"""
|
||||
embedding_provider, embedding_model_name = full_embedding_name.split('.')
|
||||
|
||||
# Calculate the embedding model to be used
|
||||
if embedding_provider == "mistral":
|
||||
api_key = current_app.config['MISTRAL_API_KEY']
|
||||
embedding_model = TrackedMistralAIEmbeddings(
|
||||
model=embedding_model_name
|
||||
)
|
||||
else:
|
||||
raise EveAIInvalidEmbeddingModel(tenant_id, catalog_id)
|
||||
|
||||
# Calculate the Embedding Model Class to be used to store embeddings
|
||||
if embedding_model_name == "mistral-embed":
|
||||
embedding_model_class = EmbeddingMistral
|
||||
else:
|
||||
raise EveAIInvalidEmbeddingModel(tenant_id, catalog_id)
|
||||
|
||||
return embedding_model, embedding_model_class
|
||||
|
||||
|
||||
def get_embedding_llm(full_model_name='mistral.mistral-small-latest', temperature=0.3):
|
||||
llm = embedding_llm_model_cache.get((full_model_name, temperature))
|
||||
if not llm:
|
||||
llm_provider, llm_model_name = full_model_name.split('.')
|
||||
if llm_provider == "openai":
|
||||
llm = ChatOpenAI(
|
||||
api_key=current_app.config['OPENAI_API_KEY'],
|
||||
model=llm_model_name,
|
||||
temperature=temperature,
|
||||
callbacks=[llm_metrics_handler]
|
||||
)
|
||||
elif llm_provider == "mistral":
|
||||
llm = ChatMistralAI(
|
||||
api_key=current_app.config['MISTRAL_API_KEY'],
|
||||
model=llm_model_name,
|
||||
temperature=temperature,
|
||||
callbacks=[llm_metrics_handler]
|
||||
)
|
||||
embedding_llm_model_cache[(full_model_name, temperature)] = llm
|
||||
|
||||
return llm
|
||||
|
||||
|
||||
def get_crewai_llm(full_model_name='mistral.mistral-large-latest', temperature=0.3):
|
||||
llm = crewai_llm_model_cache.get((full_model_name, temperature))
|
||||
if not llm:
|
||||
llm_provider, llm_model_name = full_model_name.split('.')
|
||||
crew_full_model_name = f"{llm_provider}/{llm_model_name}"
|
||||
api_key = None
|
||||
if llm_provider == "openai":
|
||||
api_key = current_app.config['OPENAI_API_KEY']
|
||||
elif llm_provider == "mistral":
|
||||
api_key = current_app.config['MISTRAL_API_KEY']
|
||||
|
||||
llm = LLM(
|
||||
model=crew_full_model_name,
|
||||
temperature=temperature,
|
||||
api_key=api_key
|
||||
)
|
||||
crewai_llm_model_cache[(full_model_name, temperature)] = llm
|
||||
|
||||
return llm
|
||||
|
||||
|
||||
def process_pdf():
|
||||
full_model_name = 'mistral-ocr-latest'
|
||||
|
||||
|
||||
def get_template(template_name: str, version: Optional[str] = "1.0", temperature: float = 0.3) -> tuple[
|
||||
Any, BaseChatModel | None | ChatOpenAI | ChatMistralAI]:
|
||||
"""
|
||||
Get a prompt template
|
||||
"""
|
||||
prompt = cache_manager.prompts_config_cache.get_config(template_name, version)
|
||||
if "llm_model" in prompt:
|
||||
llm = get_embedding_llm(full_model_name=prompt["llm_model"], temperature=temperature)
|
||||
else:
|
||||
llm = get_embedding_llm(temperature=temperature)
|
||||
|
||||
return prompt["content"], llm
|
||||
|
||||
|
||||
def get_transcription_model(model_name: str = "whisper-1") -> TrackedOpenAITranscription:
|
||||
"""
|
||||
Get a transcription model instance
|
||||
"""
|
||||
api_key = os.getenv('OPENAI_API_KEY')
|
||||
return TrackedOpenAITranscription(
|
||||
api_key=api_key,
|
||||
model=model_name
|
||||
)
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import gevent
|
||||
import time
|
||||
from flask import current_app
|
||||
@@ -28,3 +30,17 @@ def sync_folder(file_path):
|
||||
dir_fd = os.open(file_path, os.O_RDONLY)
|
||||
os.fsync(dir_fd)
|
||||
os.close(dir_fd)
|
||||
|
||||
|
||||
def get_project_root():
|
||||
"""Get the root directory of the project."""
|
||||
# Use the module that's actually running (not this file)
|
||||
module = sys.modules['__main__']
|
||||
if hasattr(module, '__file__'):
|
||||
# Get the path to the main module
|
||||
main_path = os.path.abspath(module.__file__)
|
||||
# Get the root directory (where the main module is located)
|
||||
return os.path.dirname(main_path)
|
||||
else:
|
||||
# Fallback: use current working directory
|
||||
return os.getcwd()
|
||||
|
||||
11
common/utils/prometheus_utils.py
Normal file
11
common/utils/prometheus_utils.py
Normal file
@@ -0,0 +1,11 @@
|
||||
from flask import current_app
|
||||
from prometheus_client import push_to_gateway
|
||||
|
||||
|
||||
def sanitize_label(value):
|
||||
"""Convert value to valid Prometheus label by removing/replacing invalid chars"""
|
||||
if value is None:
|
||||
return ""
|
||||
# Replace spaces and special chars with underscores
|
||||
import re
|
||||
return re.sub(r'[^a-zA-Z0-9_]', '_', str(value))
|
||||
78
common/utils/pydantic_utils.py
Normal file
78
common/utils/pydantic_utils.py
Normal file
@@ -0,0 +1,78 @@
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import Any, Dict, List, Optional, Type, Union
|
||||
|
||||
|
||||
def flatten_pydantic_model(model: BaseModel, merge_strategy: Dict[str, str] = {}) -> Dict[str, Any]:
|
||||
"""
|
||||
Flattens a nested Pydantic model by bringing all attributes to the highest level.
|
||||
|
||||
:param model: Pydantic model instance to be flattened.
|
||||
:param merge_strategy: Dictionary defining how to handle duplicate attributes.
|
||||
:return: Flattened dictionary representation of the model.
|
||||
"""
|
||||
flat_dict = {}
|
||||
|
||||
def recursive_flatten(obj: BaseModel, parent_key=""):
|
||||
for field_name, value in obj.model_dump(exclude_unset=True, by_alias=True).items():
|
||||
new_key = field_name # Maintain original field names
|
||||
|
||||
if isinstance(value, BaseModel):
|
||||
# Recursively flatten nested models
|
||||
recursive_flatten(value, new_key)
|
||||
elif isinstance(value, list) and all(isinstance(i, BaseModel) for i in value):
|
||||
# If it's a list of Pydantic models, flatten each element
|
||||
for item in value:
|
||||
recursive_flatten(item, new_key)
|
||||
else:
|
||||
if new_key in flat_dict and new_key in merge_strategy:
|
||||
# Apply merge strategy
|
||||
if merge_strategy[new_key] == "add":
|
||||
if isinstance(flat_dict[new_key], list) and isinstance(value, list):
|
||||
flat_dict[new_key] += value # Concatenate lists
|
||||
elif isinstance(flat_dict[new_key], (int, float)) and isinstance(value, (int, float)):
|
||||
flat_dict[new_key] += value # Sum numbers
|
||||
elif isinstance(flat_dict[new_key], str) and isinstance(value, str):
|
||||
flat_dict[new_key] += "\n" + value # Concatenate strings
|
||||
elif merge_strategy[new_key] == "first":
|
||||
pass # Keep the first occurrence
|
||||
elif merge_strategy[new_key] == "last":
|
||||
flat_dict[new_key] = value
|
||||
else:
|
||||
flat_dict[new_key] = value
|
||||
|
||||
recursive_flatten(model)
|
||||
return flat_dict
|
||||
|
||||
|
||||
def merge_dicts(base_dict: Dict[str, Any], new_data: Union[Dict[str, Any], BaseModel], merge_strategy: Dict[str, str]) \
|
||||
-> Dict[str, Any]:
|
||||
"""
|
||||
Merges a Pydantic model (or dictionary) into an existing dictionary based on a merge strategy.
|
||||
|
||||
:param base_dict: The base dictionary to merge into.
|
||||
:param new_data: The new Pydantic model or dictionary to merge.
|
||||
:param merge_strategy: Dict defining how to merge duplicate attributes.
|
||||
:return: Updated dictionary after merging.
|
||||
"""
|
||||
if isinstance(new_data, BaseModel):
|
||||
new_data = flatten_pydantic_model(new_data) # Convert Pydantic model to dict
|
||||
|
||||
for key, value in new_data.items():
|
||||
if key in base_dict and key in merge_strategy:
|
||||
strategy = merge_strategy[key]
|
||||
|
||||
if strategy == "add":
|
||||
if isinstance(base_dict[key], list) and isinstance(value, list):
|
||||
base_dict[key] += value # Concatenate lists
|
||||
elif isinstance(base_dict[key], (int, float)) and isinstance(value, (int, float)):
|
||||
base_dict[key] += value # Sum numbers
|
||||
elif isinstance(base_dict[key], str) and isinstance(value, str):
|
||||
base_dict[key] += " " + value # Concatenate strings
|
||||
elif strategy == "first":
|
||||
pass # Keep the first occurrence (do nothing)
|
||||
elif strategy == "last":
|
||||
base_dict[key] = value # Always overwrite with latest value
|
||||
else:
|
||||
base_dict[key] = value # Add new field
|
||||
|
||||
return base_dict
|
||||
@@ -1,19 +1,50 @@
|
||||
from flask import session, current_app
|
||||
from common.models.user import Tenant
|
||||
from sqlalchemy import and_
|
||||
|
||||
from common.models.user import Tenant, Partner
|
||||
from common.models.entitlements import License
|
||||
from common.utils.database import Database
|
||||
from common.utils.eveai_exceptions import EveAITenantNotFound, EveAITenantInvalid, EveAINoActiveLicense
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
|
||||
|
||||
# Definition of Trigger Handlers
|
||||
def set_tenant_session_data(sender, user, **kwargs):
|
||||
current_app.logger.debug(f"Setting tenant session data for user {user.id}")
|
||||
tenant = Tenant.query.filter_by(id=user.tenant_id).first()
|
||||
session['tenant'] = tenant.to_dict()
|
||||
session['default_language'] = tenant.default_language
|
||||
session['default_embedding_model'] = tenant.embedding_model
|
||||
session['default_llm_model'] = tenant.llm_model
|
||||
partner = Partner.query.filter_by(tenant_id=user.tenant_id).first()
|
||||
if partner:
|
||||
session['partner'] = partner.to_dict()
|
||||
else:
|
||||
# Remove partner from session if it exists
|
||||
session.pop('partner', None)
|
||||
|
||||
|
||||
def clear_tenant_session_data(sender, user, **kwargs):
|
||||
session.pop('tenant', None)
|
||||
session.pop('default_language', None)
|
||||
session.pop('default_embedding_model', None)
|
||||
session.pop('default_llm_model', None)
|
||||
session.pop('default_llm_model', None)
|
||||
session.pop('partner', None)
|
||||
|
||||
|
||||
def is_valid_tenant(tenant_id):
|
||||
if tenant_id == 1: # The 'root' tenant, is always valid
|
||||
return True
|
||||
tenant = Tenant.query.get(tenant_id)
|
||||
Database(tenant).switch_schema()
|
||||
if tenant is None:
|
||||
raise EveAITenantNotFound()
|
||||
elif tenant.type == 'Inactive':
|
||||
raise EveAITenantInvalid(tenant_id)
|
||||
else:
|
||||
current_date = dt.now(tz=tz.utc).date()
|
||||
# TODO -> Check vervangen door Active License Period!
|
||||
# active_license = (License.query.filter_by(tenant_id=tenant_id)
|
||||
# .filter(and_(License.start_date <= current_date,
|
||||
# License.end_date >= current_date))
|
||||
# .one_or_none())
|
||||
# if not active_license:
|
||||
# raise EveAINoActiveLicense(tenant_id)
|
||||
|
||||
return True
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
from flask import current_app, render_template
|
||||
from flask_mailman import EmailMessage
|
||||
from flask_security import current_user
|
||||
from itsdangerous import URLSafeTimedSerializer
|
||||
import socket
|
||||
|
||||
from common.models.user import Role
|
||||
from common.utils.nginx_utils import prefixed_url_for
|
||||
from common.utils.mail_utils import send_email
|
||||
|
||||
|
||||
def confirm_token(token, expiration=3600):
|
||||
@@ -11,19 +12,11 @@ def confirm_token(token, expiration=3600):
|
||||
try:
|
||||
email = serializer.loads(token, salt=current_app.config['SECURITY_PASSWORD_SALT'], max_age=expiration)
|
||||
except Exception as e:
|
||||
current_app.logger.debug(f'Error confirming token: {e}')
|
||||
current_app.logger.error(f'Error confirming token: {e}')
|
||||
raise
|
||||
return email
|
||||
|
||||
|
||||
def send_email(to, subject, template):
|
||||
msg = EmailMessage(subject=subject,
|
||||
body=template,
|
||||
to=[to])
|
||||
msg.content_subtype = "html"
|
||||
msg.send()
|
||||
|
||||
|
||||
def generate_reset_token(email):
|
||||
serializer = URLSafeTimedSerializer(current_app.config['SECRET_KEY'])
|
||||
return serializer.dumps(email, salt=current_app.config['SECURITY_PASSWORD_SALT'])
|
||||
@@ -35,20 +28,14 @@ def generate_confirmation_token(email):
|
||||
|
||||
|
||||
def send_confirmation_email(user):
|
||||
current_app.logger.debug(f'Sending confirmation email to {user.email}')
|
||||
|
||||
if not test_smtp_connection():
|
||||
raise Exception("Failed to connect to SMTP server")
|
||||
|
||||
token = generate_confirmation_token(user.email)
|
||||
confirm_url = prefixed_url_for('security_bp.confirm_email', token=token, _external=True)
|
||||
current_app.logger.debug(f'Confirmation URL: {confirm_url}')
|
||||
|
||||
html = render_template('email/activate.html', confirm_url=confirm_url)
|
||||
subject = "Please confirm your email"
|
||||
|
||||
try:
|
||||
send_email(user.email, "Confirm your email", html)
|
||||
send_email(user.email, f"{user.first_name} {user.last_name}", "Confirm your email", html)
|
||||
current_app.logger.info(f'Confirmation email sent to {user.email}')
|
||||
except Exception as e:
|
||||
current_app.logger.error(f'Failed to send confirmation email to {user.email}. Error: {str(e)}')
|
||||
@@ -56,46 +43,56 @@ def send_confirmation_email(user):
|
||||
|
||||
|
||||
def send_reset_email(user):
|
||||
current_app.logger.debug(f'Sending reset email to {user.email}')
|
||||
token = generate_reset_token(user.email)
|
||||
reset_url = prefixed_url_for('security_bp.reset_password', token=token, _external=True)
|
||||
current_app.logger.debug(f'Reset URL: {reset_url}')
|
||||
|
||||
html = render_template('email/reset_password.html', reset_url=reset_url)
|
||||
subject = "Reset Your Password"
|
||||
|
||||
try:
|
||||
send_email(user.email, "Reset Your Password", html)
|
||||
send_email(user.email, f"{user.first_name} {user.last_name}", subject, html)
|
||||
current_app.logger.info(f'Reset email sent to {user.email}')
|
||||
except Exception as e:
|
||||
current_app.logger.error(f'Failed to send reset email to {user.email}. Error: {str(e)}')
|
||||
raise
|
||||
|
||||
|
||||
def test_smtp_connection():
|
||||
try:
|
||||
current_app.logger.info(f"Attempting to resolve google.com...")
|
||||
google_ip = socket.gethostbyname('google.com')
|
||||
current_app.logger.info(f"Successfully resolved google.com to {google_ip}")
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Failed to resolve google.com: {str(e)}")
|
||||
def get_current_user_roles():
|
||||
"""Get the roles of the currently authenticated user.
|
||||
|
||||
try:
|
||||
smtp_server = current_app.config['MAIL_SERVER']
|
||||
current_app.logger.info(f"Attempting to resolve {smtp_server}...")
|
||||
smtp_ip = socket.gethostbyname(smtp_server)
|
||||
current_app.logger.info(f"Successfully resolved {smtp_server} to {smtp_ip}")
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Failed to resolve {smtp_server}: {str(e)}")
|
||||
Returns:
|
||||
List of Role objects or empty list if no user is authenticated
|
||||
"""
|
||||
if current_user.is_authenticated:
|
||||
return current_user.roles
|
||||
return []
|
||||
|
||||
try:
|
||||
smtp_server = current_app.config['MAIL_SERVER']
|
||||
smtp_port = current_app.config['MAIL_PORT']
|
||||
sock = socket.create_connection((smtp_server, smtp_port), timeout=10)
|
||||
sock.close()
|
||||
current_app.logger.info(f"Successfully connected to SMTP server {smtp_server}:{smtp_port}")
|
||||
return True
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Failed to connect to SMTP server: {str(e)}")
|
||||
|
||||
def current_user_has_role(role_name):
|
||||
"""Check if the current user has the specified role.
|
||||
|
||||
Args:
|
||||
role_name (str): Name of the role to check
|
||||
|
||||
Returns:
|
||||
bool: True if user has the role, False otherwise
|
||||
"""
|
||||
if not current_user.is_authenticated:
|
||||
return False
|
||||
|
||||
return any(role.name == role_name for role in current_user.roles)
|
||||
|
||||
|
||||
def current_user_roles():
|
||||
"""Get the roles of the currently authenticated user.
|
||||
|
||||
Returns:
|
||||
List of Role objects or empty list if no user is authenticated
|
||||
"""
|
||||
if current_user.is_authenticated:
|
||||
return current_user.roles
|
||||
return []
|
||||
|
||||
|
||||
def all_user_roles():
|
||||
roles = [(role.id, role.name) for role in Role.query.all()]
|
||||
|
||||
@@ -4,7 +4,7 @@ from flask import Flask
|
||||
|
||||
|
||||
def generate_api_key(prefix="EveAI-Chat"):
|
||||
parts = [str(random.randint(1000, 9999)) for _ in range(5)]
|
||||
parts = [str(random.randint(1000, 9999)) for _ in range(8)]
|
||||
return f"{prefix}-{'-'.join(parts)}"
|
||||
|
||||
|
||||
|
||||
196
common/utils/specialist_utils.py
Normal file
196
common/utils/specialist_utils.py
Normal file
@@ -0,0 +1,196 @@
|
||||
from datetime import datetime as dt, timezone as tz
|
||||
from typing import Optional, Dict, Any
|
||||
from flask import current_app
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from common.extensions import db, cache_manager
|
||||
from common.models.interaction import (
|
||||
Specialist, EveAIAgent, EveAITask, EveAITool
|
||||
)
|
||||
from common.utils.model_logging_utils import set_logging_information, update_logging_information
|
||||
|
||||
|
||||
def initialize_specialist(specialist_id: int, specialist_type: str, specialist_version: str):
|
||||
"""
|
||||
Initialize an agentic specialist by creating all its components based on configuration.
|
||||
|
||||
Args:
|
||||
specialist_id: ID of the specialist to initialize
|
||||
specialist_type: Type of the specialist
|
||||
specialist_version: Version of the specialist type to use
|
||||
|
||||
Raises:
|
||||
ValueError: If specialist not found or invalid configuration
|
||||
SQLAlchemyError: If database operations fail
|
||||
"""
|
||||
config = cache_manager.specialists_config_cache.get_config(specialist_type, specialist_version)
|
||||
if not config:
|
||||
raise ValueError(f"No configuration found for {specialist_type} version {specialist_version}")
|
||||
if config['framework'] == 'langchain':
|
||||
pass # Langchain does not require additional items to be initialized. All configuration is in the specialist.
|
||||
|
||||
specialist = Specialist.query.get(specialist_id)
|
||||
if not specialist:
|
||||
raise ValueError(f"Specialist with ID {specialist_id} not found")
|
||||
|
||||
if config['framework'] == 'crewai':
|
||||
initialize_crewai_specialist(specialist, config)
|
||||
|
||||
|
||||
def initialize_crewai_specialist(specialist: Specialist, config: Dict[str, Any]):
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
try:
|
||||
# Initialize agents
|
||||
if 'agents' in config:
|
||||
for agent_config in config['agents']:
|
||||
_create_agent(
|
||||
specialist_id=specialist.id,
|
||||
agent_type=agent_config['type'],
|
||||
agent_version=agent_config['version'],
|
||||
name=agent_config.get('name'),
|
||||
description=agent_config.get('description'),
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
# Initialize tasks
|
||||
if 'tasks' in config:
|
||||
for task_config in config['tasks']:
|
||||
_create_task(
|
||||
specialist_id=specialist.id,
|
||||
task_type=task_config['type'],
|
||||
task_version=task_config['version'],
|
||||
name=task_config.get('name'),
|
||||
description=task_config.get('description'),
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
# Initialize tools
|
||||
if 'tools' in config:
|
||||
for tool_config in config['tools']:
|
||||
_create_tool(
|
||||
specialist_id=specialist.id,
|
||||
tool_type=tool_config['type'],
|
||||
tool_version=tool_config['version'],
|
||||
name=tool_config.get('name'),
|
||||
description=tool_config.get('description'),
|
||||
timestamp=timestamp
|
||||
)
|
||||
|
||||
db.session.commit()
|
||||
current_app.logger.info(f"Successfully initialized crewai specialist {specialist.id}")
|
||||
|
||||
except SQLAlchemyError as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Database error initializing crewai specialist {specialist.id}: {str(e)}")
|
||||
raise
|
||||
except Exception as e:
|
||||
db.session.rollback()
|
||||
current_app.logger.error(f"Error initializing crewai specialist {specialist.id}: {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def _create_agent(
|
||||
specialist_id: int,
|
||||
agent_type: str,
|
||||
agent_version: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
timestamp: Optional[dt] = None
|
||||
) -> EveAIAgent:
|
||||
"""Create an agent with the given configuration."""
|
||||
if timestamp is None:
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
# Get agent configuration from cache
|
||||
agent_config = cache_manager.agents_config_cache.get_config(agent_type, agent_version)
|
||||
|
||||
agent = EveAIAgent(
|
||||
specialist_id=specialist_id,
|
||||
name=name or agent_config.get('name', agent_type),
|
||||
description=description or agent_config.get('metadata').get('description', ''),
|
||||
type=agent_type,
|
||||
type_version=agent_version,
|
||||
role=None,
|
||||
goal=None,
|
||||
backstory=None,
|
||||
tuning=False,
|
||||
configuration=None,
|
||||
arguments=None
|
||||
)
|
||||
|
||||
set_logging_information(agent, timestamp)
|
||||
|
||||
db.session.add(agent)
|
||||
current_app.logger.info(f"Created agent {agent.id} of type {agent_type}")
|
||||
return agent
|
||||
|
||||
|
||||
def _create_task(
|
||||
specialist_id: int,
|
||||
task_type: str,
|
||||
task_version: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
timestamp: Optional[dt] = None
|
||||
) -> EveAITask:
|
||||
"""Create a task with the given configuration."""
|
||||
if timestamp is None:
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
# Get task configuration from cache
|
||||
task_config = cache_manager.tasks_config_cache.get_config(task_type, task_version)
|
||||
|
||||
task = EveAITask(
|
||||
specialist_id=specialist_id,
|
||||
name=name or task_config.get('name', task_type),
|
||||
description=description or task_config.get('metadata').get('description', ''),
|
||||
type=task_type,
|
||||
type_version=task_version,
|
||||
task_description=None,
|
||||
expected_output=None,
|
||||
tuning=False,
|
||||
configuration=None,
|
||||
arguments=None,
|
||||
context=None,
|
||||
asynchronous=False,
|
||||
)
|
||||
|
||||
set_logging_information(task, timestamp)
|
||||
|
||||
db.session.add(task)
|
||||
current_app.logger.info(f"Created task {task.id} of type {task_type}")
|
||||
return task
|
||||
|
||||
|
||||
def _create_tool(
|
||||
specialist_id: int,
|
||||
tool_type: str,
|
||||
tool_version: str,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
timestamp: Optional[dt] = None
|
||||
) -> EveAITool:
|
||||
"""Create a tool with the given configuration."""
|
||||
if timestamp is None:
|
||||
timestamp = dt.now(tz=tz.utc)
|
||||
|
||||
# Get tool configuration from cache
|
||||
tool_config = cache_manager.tools_config_cache.get_config(tool_type, tool_version)
|
||||
|
||||
tool = EveAITool(
|
||||
specialist_id=specialist_id,
|
||||
name=name or tool_config.get('name', tool_type),
|
||||
description=description or tool_config.get('metadata').get('description', ''),
|
||||
type=tool_type,
|
||||
type_version=tool_version,
|
||||
tuning=False,
|
||||
configuration=None,
|
||||
arguments=None,
|
||||
)
|
||||
|
||||
set_logging_information(tool, timestamp)
|
||||
|
||||
db.session.add(tool)
|
||||
current_app.logger.info(f"Created tool {tool.id} of type {tool_type}")
|
||||
return tool
|
||||
48
common/utils/startup_eveai.py
Normal file
48
common/utils/startup_eveai.py
Normal file
@@ -0,0 +1,48 @@
|
||||
import time
|
||||
|
||||
from redis import Redis
|
||||
|
||||
from common.extensions import cache_manager
|
||||
|
||||
|
||||
def perform_startup_actions(app):
|
||||
perform_startup_invalidation(app)
|
||||
|
||||
|
||||
def perform_startup_invalidation(app):
|
||||
"""
|
||||
Perform cache invalidation only once during startup using a persistent marker (also called flag or semaphore
|
||||
- see docs).
|
||||
Uses a combination of lock and marker to ensure invalidation happens exactly once
|
||||
per deployment.
|
||||
"""
|
||||
redis_client = Redis.from_url(app.config['REDIS_BASE_URI'])
|
||||
startup_time = int(time.time())
|
||||
marker_key = 'startup_invalidation_completed'
|
||||
lock_key = 'startup_invalidation_lock'
|
||||
|
||||
try:
|
||||
# First try to get the lock
|
||||
lock = redis_client.lock(lock_key, timeout=30)
|
||||
if lock.acquire(blocking=False):
|
||||
try:
|
||||
# Check if invalidation was already performed
|
||||
if not redis_client.get(marker_key):
|
||||
# Perform invalidation
|
||||
cache_manager.invalidate_region('eveai_config')
|
||||
cache_manager.invalidate_region('eveai_chat_workers')
|
||||
|
||||
redis_client.setex(marker_key, 180, str(startup_time))
|
||||
app.logger.info("Startup cache invalidation completed")
|
||||
else:
|
||||
app.logger.info("Startup cache invalidation already performed")
|
||||
finally:
|
||||
lock.release()
|
||||
else:
|
||||
app.logger.info("Another process is handling startup invalidation")
|
||||
|
||||
except Exception as e:
|
||||
app.logger.error(f"Error during startup invalidation: {e}")
|
||||
# In case of error, we don't want to block the application startup
|
||||
pass
|
||||
|
||||
112
common/utils/string_list_converter.py
Normal file
112
common/utils/string_list_converter.py
Normal file
@@ -0,0 +1,112 @@
|
||||
from typing import List, Union
|
||||
import re
|
||||
|
||||
|
||||
class StringListConverter:
|
||||
"""Utility class for converting between comma-separated strings and lists"""
|
||||
|
||||
@staticmethod
|
||||
def string_to_list(input_string: Union[str, None], allow_empty: bool = True) -> List[str]:
|
||||
"""
|
||||
Convert a comma-separated string to a list of strings.
|
||||
|
||||
Args:
|
||||
input_string: Comma-separated string to convert
|
||||
allow_empty: If True, returns empty list for None/empty input
|
||||
If False, raises ValueError for None/empty input
|
||||
|
||||
Returns:
|
||||
List of stripped strings
|
||||
|
||||
Raises:
|
||||
ValueError: If input is None/empty and allow_empty is False
|
||||
"""
|
||||
if not input_string:
|
||||
if allow_empty:
|
||||
return []
|
||||
raise ValueError("Input string cannot be None or empty")
|
||||
|
||||
return [item.strip() for item in input_string.split(',') if item.strip()]
|
||||
|
||||
@staticmethod
|
||||
def list_to_string(input_list: Union[List[str], None], allow_empty: bool = True) -> str:
|
||||
"""
|
||||
Convert a list of strings to a comma-separated string.
|
||||
|
||||
Args:
|
||||
input_list: List of strings to convert
|
||||
allow_empty: If True, returns empty string for None/empty input
|
||||
If False, raises ValueError for None/empty input
|
||||
|
||||
Returns:
|
||||
Comma-separated string
|
||||
|
||||
Raises:
|
||||
ValueError: If input is None/empty and allow_empty is False
|
||||
"""
|
||||
if not input_list:
|
||||
if allow_empty:
|
||||
return ''
|
||||
raise ValueError("Input list cannot be None or empty")
|
||||
|
||||
return ', '.join(str(item).strip() for item in input_list)
|
||||
|
||||
@staticmethod
|
||||
def validate_format(input_string: str,
|
||||
allowed_chars: str = r'a-zA-Z0-9_\-',
|
||||
min_length: int = 1,
|
||||
max_length: int = 50) -> bool:
|
||||
"""
|
||||
Validate the format of items in a comma-separated string.
|
||||
|
||||
Args:
|
||||
input_string: String to validate
|
||||
allowed_chars: String of allowed characters (for regex pattern)
|
||||
min_length: Minimum length for each item
|
||||
max_length: Maximum length for each item
|
||||
|
||||
Returns:
|
||||
bool: True if format is valid, False otherwise
|
||||
"""
|
||||
if not input_string:
|
||||
return False
|
||||
|
||||
# Create regex pattern for individual items
|
||||
pattern = f'^[{allowed_chars}]{{{min_length},{max_length}}}$'
|
||||
|
||||
try:
|
||||
# Convert to list and check each item
|
||||
items = StringListConverter.string_to_list(input_string)
|
||||
return all(bool(re.match(pattern, item)) for item in items)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def validate_and_convert(input_string: str,
|
||||
allowed_chars: str = r'a-zA-Z0-9_\-',
|
||||
min_length: int = 1,
|
||||
max_length: int = 50) -> List[str]:
|
||||
"""
|
||||
Validate and convert a comma-separated string to a list.
|
||||
|
||||
Args:
|
||||
input_string: String to validate and convert
|
||||
allowed_chars: String of allowed characters (for regex pattern)
|
||||
min_length: Minimum length for each item
|
||||
max_length: Maximum length for each item
|
||||
|
||||
Returns:
|
||||
List of validated and converted strings
|
||||
|
||||
Raises:
|
||||
ValueError: If input string format is invalid
|
||||
"""
|
||||
if not StringListConverter.validate_format(
|
||||
input_string, allowed_chars, min_length, max_length
|
||||
):
|
||||
raise ValueError(
|
||||
f"Invalid format. Items must be {min_length}-{max_length} characters "
|
||||
f"long and contain only these characters: {allowed_chars}"
|
||||
)
|
||||
|
||||
return StringListConverter.string_to_list(input_string)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user