299 lines
13 KiB
Python
299 lines
13 KiB
Python
# Part of Odoo. See LICENSE file for full copyright and licensing details.
|
|
import logging
|
|
from collections import defaultdict
|
|
|
|
from odoo import api, fields, models, _
|
|
from requests.exceptions import RequestException
|
|
from odoo.tools import SQL
|
|
from odoo.exceptions import UserError
|
|
|
|
from odoo.addons.ai.orm.field_vector import Vector
|
|
from odoo.addons.ai.utils.llm_api_service import LLMApiService
|
|
from odoo.addons.ai.utils.llm_providers import (
|
|
EMBEDDING_MODELS_SELECTION,
|
|
get_embedding_config,
|
|
get_provider_for_embedding_model,
|
|
get_embedding_model,
|
|
)
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AIEmbedding(models.Model):
|
|
_name = 'ai.embedding'
|
|
_description = "Attachment Chunks Embedding"
|
|
_order = 'sequence'
|
|
|
|
attachment_id = fields.Many2one(
|
|
'ir.attachment',
|
|
string="Attachment",
|
|
required=True,
|
|
ondelete='cascade'
|
|
)
|
|
checksum = fields.Char(related='attachment_id.checksum')
|
|
sequence = fields.Integer(string="Sequence", default=10)
|
|
content = fields.Text(string="Chunk Content", required=True)
|
|
embedding_model = fields.Selection(selection=EMBEDDING_MODELS_SELECTION, string="Embedding Model", required=True)
|
|
has_embedding_generation_failed = fields.Boolean(string="Has Embedding Generation Failed", default=False)
|
|
embedding_vector = Vector(size=1536)
|
|
_embedding_vector_idx = models.Index("USING ivfflat (embedding_vector vector_cosine_ops)")
|
|
|
|
@api.model
|
|
def _get_dimensions(self):
|
|
return self._fields['embedding_vector'].size
|
|
|
|
@api.model
|
|
def _get_similar_chunks(self, query_embedding, sources, embedding_model, top_n=5):
|
|
active_sources = sources.filtered(lambda s: s.is_active)
|
|
if not active_sources:
|
|
return self
|
|
|
|
attachment_ids = active_sources.mapped('attachment_id').ids
|
|
target_checksums = self.env['ir.attachment'].browse(attachment_ids).mapped('checksum')
|
|
# Execute the SQL query to find similar embeddings of the same sources' attachments checksum
|
|
return self.browse(id_ for id_, *_ in self.env.execute_query(SQL(
|
|
'''
|
|
SELECT
|
|
ai_embedding.id,
|
|
1 - (embedding_vector <=> %s::vector) AS similarity
|
|
FROM ai_embedding
|
|
INNER JOIN ir_attachment ON ir_attachment.id = ai_embedding.attachment_id
|
|
WHERE ir_attachment.checksum = ANY(%s) AND ai_embedding.embedding_model = %s
|
|
ORDER BY similarity DESC
|
|
LIMIT %s;
|
|
''',
|
|
query_embedding, target_checksums, embedding_model, top_n)
|
|
)
|
|
)
|
|
|
|
@api.model
|
|
def _cron_generate_embedding(self, batch_size=100):
|
|
"""
|
|
Generate embeddings for sources, handling multiple embedding models per source.
|
|
"""
|
|
# Check for sources that need to be chunked and create embedding chunks
|
|
self._create_embedding_chunks()
|
|
|
|
# Check for all the embedding vectors to be generated
|
|
missing_embeddings = self.search([
|
|
("embedding_vector", "=", False),
|
|
("has_embedding_generation_failed", "=", False)
|
|
])
|
|
|
|
if not missing_embeddings:
|
|
return False
|
|
missing_embeddings_batch = missing_embeddings[:batch_size]
|
|
|
|
# Group by embedding model to batch similar requests
|
|
embeddings_by_model = defaultdict(list)
|
|
for embedding in missing_embeddings_batch:
|
|
model = embedding.embedding_model
|
|
embeddings_by_model[model].append(embedding)
|
|
|
|
staged_sources = self.env['ai.agent.source'].search([
|
|
('attachment_id.checksum', 'in', missing_embeddings_batch.mapped('checksum'))
|
|
])
|
|
_logger.info("Starting embedding update - missing %s embeddings.", len(missing_embeddings_batch))
|
|
self.env['ir.cron']._commit_progress(remaining=len(missing_embeddings_batch))
|
|
failed_embeddings = self.env[self._name]
|
|
|
|
# Process each model group separately
|
|
for model, embeddings in embeddings_by_model.items():
|
|
provider = get_provider_for_embedding_model(self.env, model)
|
|
# Create batches respecting provider limits
|
|
batches = self._create_batches(embeddings, provider)
|
|
_logger.info(
|
|
"Processing %s embeddings for model %s in %s batches",
|
|
len(embeddings), model, len(batches)
|
|
)
|
|
|
|
for batch_idx, batch in enumerate(batches):
|
|
try:
|
|
_logger.info(
|
|
"Processing batch %s/%s with %s embeddings for model %s",
|
|
batch_idx + 1, len(batches), len(batch), model
|
|
)
|
|
|
|
# Prepare batch input
|
|
llm_service = LLMApiService(env=self.env, provider=provider)
|
|
batch_content = [
|
|
llm_service._format_for_embedding(content=emb.content, mode="document", title=emb.attachment_id.name)
|
|
for emb in batch
|
|
]
|
|
# Get embeddings for the entire batch
|
|
response = llm_service.get_embedding(
|
|
input=batch_content,
|
|
dimensions=self._get_dimensions(),
|
|
model=get_embedding_model(provider),
|
|
)
|
|
# Update each embedding record with its corresponding vector
|
|
for idx, embedding in enumerate(batch):
|
|
try:
|
|
embedding.embedding_vector = response['data'][idx]['embedding']
|
|
except KeyError as e:
|
|
_logger.error(
|
|
"Failed to extract embedding for record %s: %s",
|
|
embedding.id, str(e)
|
|
)
|
|
failed_embeddings |= embedding
|
|
|
|
# Commit progress for this batch
|
|
if not self.env['ir.cron']._commit_progress(len(batch)):
|
|
break
|
|
|
|
except (RequestException, UserError, AttributeError) as e:
|
|
_logger.error(
|
|
"Failed to process batch %s/%s for model %s: %s",
|
|
batch_idx + 1, len(batches), model, str(e)
|
|
)
|
|
# Mark all embeddings in failed batch as failed
|
|
failed_embeddings |= self.env[self._name].concat(*batch)
|
|
continue
|
|
|
|
if failed_embeddings:
|
|
failed_embeddings.write({'has_embedding_generation_failed': True})
|
|
self.env['ir.cron']._commit_progress(len(failed_embeddings))
|
|
|
|
# Handle the status of sources
|
|
for source in staged_sources:
|
|
embedding_model = source.agent_id._get_embedding_model()
|
|
source._update_source_status(embedding_model)
|
|
|
|
if len(missing_embeddings) > batch_size:
|
|
# we still have unfinished embeddings to generate: run the CRON again
|
|
self.env.ref('ai.ir_cron_generate_embedding')._trigger()
|
|
return True
|
|
|
|
return False
|
|
|
|
@api.autovacuum
|
|
def _cron_remove_deprecated_models_embeddings_and_recompute(self):
|
|
embeddings_to_unlink = self.env['ai.embedding'].search(
|
|
[('embedding_model', 'not in', [selection[0] for selection in EMBEDDING_MODELS_SELECTION])]
|
|
)
|
|
checksum_providers_to_embed = {
|
|
(embedding.checksum, get_provider_for_embedding_model(self.env, embedding.embedding_model))
|
|
for embedding in embeddings_to_unlink
|
|
}
|
|
embeddings_to_unlink.unlink()
|
|
|
|
candidate_sources = self.env['ai.agent.source'].search([
|
|
('attachment_id.checksum', 'in', [checksum for checksum, _ in checksum_providers_to_embed])
|
|
])
|
|
# The same attachment can be embedded multiple times using models from different providers.
|
|
# A source linked to an attachment will be re-embed only if the embedding model of the provider
|
|
# of its agent has been deprecated.
|
|
sources_to_embed = candidate_sources.filtered(
|
|
lambda source: (source.attachment_id.checksum, source.agent_id._get_provider()) in checksum_providers_to_embed
|
|
)
|
|
sources_to_embed.write({
|
|
'status': 'processing',
|
|
'is_active': False,
|
|
'error_details': False,
|
|
})
|
|
self.env.ref('ai.ir_cron_generate_embedding')._trigger()
|
|
|
|
def _create_embedding_chunks(self):
|
|
"""
|
|
Create embedding chunks for sources that are processing and have an attachment
|
|
and that don't have any embedding chunks yet
|
|
"""
|
|
sources_to_process = self.env['ai.agent.source'].search([
|
|
('status', '=', 'processing'),
|
|
('attachment_id', '!=', False),
|
|
])
|
|
if sources_to_process:
|
|
existing_embeddings_grouped = self._read_group(
|
|
domain=[],
|
|
groupby=['checksum', 'embedding_model'],
|
|
)
|
|
existing_checksum_model_pairs = [(checksum, embedding_model) for checksum, embedding_model in existing_embeddings_grouped]
|
|
for source in sources_to_process:
|
|
embedding_model = source.agent_id._get_embedding_model()
|
|
if (source.attachment_id.checksum, embedding_model) not in existing_checksum_model_pairs:
|
|
_logger.info("Creating embedding chunks for source %s", source.name)
|
|
content = source.attachment_id._get_attachment_content()
|
|
if not content:
|
|
source.write({
|
|
'status': 'failed',
|
|
'error_details': _("Invalid attachment. Failed to extract content."),
|
|
})
|
|
continue
|
|
source.attachment_id._setup_attachment_chunks(embedding_model, content)
|
|
existing_checksum_model_pairs.append((source.attachment_id.checksum, embedding_model))
|
|
|
|
def _estimate_tokens(self, text):
|
|
"""Estimate token count based on text length.
|
|
Based on https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them
|
|
:param text: Text to estimate tokens for
|
|
:type text: str
|
|
:return: Estimated token count
|
|
:rtype: int
|
|
"""
|
|
return len(text) // 4 if text else 0
|
|
|
|
def _create_batches(self, embeddings, provider):
|
|
"""
|
|
Group embeddings into batches respecting provider limits.
|
|
:param embeddings: List of embeddings to group into batches
|
|
:type embeddings: list[AIEmbedding]
|
|
:param provider: Provider to use for embedding generation
|
|
:type provider: str
|
|
:return: List of batches
|
|
:rtype: list[list[AIEmbedding]]
|
|
"""
|
|
config = get_embedding_config(self.env, provider)
|
|
max_batch_size = config['max_batch_size']
|
|
max_tokens = config['max_tokens_per_request']
|
|
|
|
batches = []
|
|
current_batch = []
|
|
current_batch_tokens = 0
|
|
|
|
for embedding in embeddings:
|
|
content_tokens = self._estimate_tokens(embedding.content)
|
|
if (len(current_batch) >= max_batch_size) or \
|
|
(current_batch_tokens + content_tokens > max_tokens):
|
|
batches.append(current_batch)
|
|
current_batch = []
|
|
current_batch_tokens = 0
|
|
|
|
current_batch.append(embedding)
|
|
current_batch_tokens += content_tokens
|
|
|
|
if current_batch:
|
|
batches.append(current_batch)
|
|
|
|
return batches
|
|
|
|
def _get_indexed_embedding_models_by_checksum(self, checksum):
|
|
"""
|
|
Get indexed embedding models by checksum.
|
|
:param checksum: checksum of the attachment
|
|
:type checksum: str
|
|
:return: list of indexed embedding models of the input checksum
|
|
:rtype: list of str
|
|
"""
|
|
embeddings_by_model = self._read_group(
|
|
domain=[('checksum', '=', checksum)],
|
|
groupby=['embedding_model'],
|
|
)
|
|
indexed_embedding_models = [embedding_model[0] for embedding_model in embeddings_by_model]
|
|
return indexed_embedding_models
|
|
|
|
@api.autovacuum
|
|
def _gc_embeddings(self):
|
|
"""
|
|
Autovacuum: Cleanup embedding chunks not associated with any agent's attachments.
|
|
"""
|
|
all_agents = self.env['ai.agent'].with_context(active_test=False).search([])
|
|
used_attachment_ids = all_agents.mapped('sources_ids').mapped('attachment_id')
|
|
used_checksums = used_attachment_ids.mapped('checksum')
|
|
unused_chunks = self.search([
|
|
('checksum', 'not in', used_checksums),
|
|
])
|
|
if unused_chunks:
|
|
chunk_count = len(unused_chunks)
|
|
_logger.info("Autovacuum: Cleaning up %s unused embedding chunks", chunk_count)
|
|
unused_chunks.unlink()
|