197 lines
7.4 KiB
Python
197 lines
7.4 KiB
Python
|
"""Ontology selection agent for OntoCast.
|
||
|
|
|
||
|
|
This module provides functionality for selecting appropriate ontologies based on
|
||
|
|
the content of source text segments, ensuring that the chosen ontology best matches the
|
||
|
|
domain and requirements of the text.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
from langchain_core.output_parsers import PydanticOutputParser
|
||
|
|
from langchain_core.prompts import PromptTemplate
|
||
|
|
|
||
|
|
from ontocast.agent.common import call_llm_with_retry
|
||
|
|
from ontocast.onto.enum import Status
|
||
|
|
from ontocast.onto.model import create_ontology_selector_report_model
|
||
|
|
from ontocast.onto.null import NULL_ONTOLOGY
|
||
|
|
from ontocast.onto.state import AgentState
|
||
|
|
from ontocast.prompt.select_ontology import template_prompt
|
||
|
|
from ontocast.tool import OntologyManager
|
||
|
|
from ontocast.toolbox import ToolBox
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def _create_document_excerpt(state: AgentState, max_length: int = 3000) -> str:
|
||
|
|
"""Create a representative excerpt from the document for ontology selection.
|
||
|
|
|
||
|
|
This function samples text from multiple content units to provide a better
|
||
|
|
representation of the document content than just the first unit.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
state: The current agent state.
|
||
|
|
max_length: Maximum total length of the excerpt.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
str: A representative excerpt from the document.
|
||
|
|
"""
|
||
|
|
excerpt_parts = []
|
||
|
|
total_length = 0
|
||
|
|
chunk_length = max_length // 3 # Aim for ~3 source chunks, ~1000 chars each
|
||
|
|
|
||
|
|
# Strategy: Sample from first, middle, and last units if available
|
||
|
|
if state.content_units:
|
||
|
|
num_chunks = len(state.content_units)
|
||
|
|
indices_to_sample = []
|
||
|
|
|
||
|
|
if num_chunks == 1:
|
||
|
|
indices_to_sample = [0]
|
||
|
|
elif num_chunks == 2:
|
||
|
|
indices_to_sample = [0, 1]
|
||
|
|
else:
|
||
|
|
# Sample first, middle, and last
|
||
|
|
indices_to_sample = [0, num_chunks // 2, num_chunks - 1]
|
||
|
|
|
||
|
|
for idx in indices_to_sample:
|
||
|
|
if idx < num_chunks and total_length < max_length:
|
||
|
|
chunk_text = state.content_units[idx].text
|
||
|
|
# Take a portion of this source chunk
|
||
|
|
remaining = max_length - total_length
|
||
|
|
sample_length = min(chunk_length, remaining, len(chunk_text))
|
||
|
|
|
||
|
|
if sample_length > 0:
|
||
|
|
if sample_length < len(chunk_text):
|
||
|
|
excerpt_parts.append(chunk_text[:sample_length] + " ...")
|
||
|
|
else:
|
||
|
|
excerpt_parts.append(chunk_text)
|
||
|
|
total_length += sample_length
|
||
|
|
|
||
|
|
if excerpt_parts:
|
||
|
|
return "\n\n[...]\n\n".join(excerpt_parts)
|
||
|
|
|
||
|
|
# Fallback: Use input_text if available
|
||
|
|
if state.input_text:
|
||
|
|
if len(state.input_text) <= max_length:
|
||
|
|
return state.input_text
|
||
|
|
return state.input_text[:max_length] + " ..."
|
||
|
|
|
||
|
|
# Last resort: Use current content unit
|
||
|
|
if state.current_content_unit and state.current_content_unit.text:
|
||
|
|
chunk_text = state.current_content_unit.text
|
||
|
|
if len(chunk_text) <= max_length:
|
||
|
|
return chunk_text
|
||
|
|
return chunk_text[:max_length] + " ..."
|
||
|
|
|
||
|
|
return ""
|
||
|
|
|
||
|
|
|
||
|
|
async def select_ontology(state: AgentState, tools: ToolBox) -> AgentState:
|
||
|
|
"""Select an appropriate ontology for the document.
|
||
|
|
|
||
|
|
This function analyzes the document and selects the most appropriate
|
||
|
|
ontology based on its content and requirements using a numbered list selection.
|
||
|
|
If an ontology is already selected, it skips selection to ensure one ontology
|
||
|
|
per document.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
state: The current agent state containing the document to process.
|
||
|
|
tools: The toolbox instance providing utility functions.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
AgentState: Updated state with selected ontology.
|
||
|
|
"""
|
||
|
|
# Skip if ontology already selected (for subsequent chunks in the loop)
|
||
|
|
if not state.current_ontology.is_null():
|
||
|
|
logger.debug(
|
||
|
|
f"Ontology already selected: {state.current_ontology.ontology_id}, "
|
||
|
|
"skipping selection to maintain one ontology per document"
|
||
|
|
)
|
||
|
|
state.status = Status.SUCCESS
|
||
|
|
return state
|
||
|
|
|
||
|
|
progress_info = state.get_content_unit_progress_string()
|
||
|
|
logger.info(f"Selecting ontology for document ({progress_info})")
|
||
|
|
llm_tool = tools.llm
|
||
|
|
om_tool: OntologyManager = tools.ontology_manager
|
||
|
|
|
||
|
|
if om_tool.has_ontologies:
|
||
|
|
ontologies = om_tool.ontologies
|
||
|
|
num_ontologies = len(ontologies)
|
||
|
|
|
||
|
|
# Create numbered list of ontologies
|
||
|
|
ontologies_list_lines = []
|
||
|
|
for i, ontology in enumerate(ontologies, start=1):
|
||
|
|
ontologies_list_lines.append(f"{i}. {ontology.describe()}")
|
||
|
|
|
||
|
|
ontologies_list = "\n\n".join(ontologies_list_lines)
|
||
|
|
|
||
|
|
logger.info(f"Presenting {num_ontologies} ontologies for selection")
|
||
|
|
|
||
|
|
# Create a better document excerpt using multiple chunks
|
||
|
|
excerpt = _create_document_excerpt(state, max_length=3000)
|
||
|
|
|
||
|
|
# Create dynamic model with correct constraint
|
||
|
|
ontology_selector_report_model = create_ontology_selector_report_model(
|
||
|
|
num_ontologies
|
||
|
|
)
|
||
|
|
parser = PydanticOutputParser(pydantic_object=ontology_selector_report_model)
|
||
|
|
|
||
|
|
prompt = PromptTemplate(
|
||
|
|
template=template_prompt,
|
||
|
|
input_variables=[
|
||
|
|
"excerpt",
|
||
|
|
"ontologies_list",
|
||
|
|
"num_ontologies",
|
||
|
|
"format_instructions",
|
||
|
|
],
|
||
|
|
)
|
||
|
|
|
||
|
|
selector = await call_llm_with_retry(
|
||
|
|
llm_tool=llm_tool,
|
||
|
|
prompt=prompt,
|
||
|
|
parser=parser,
|
||
|
|
prompt_kwargs={
|
||
|
|
"excerpt": excerpt,
|
||
|
|
"ontologies_list": ontologies_list,
|
||
|
|
"num_ontologies": num_ontologies,
|
||
|
|
"format_instructions": parser.get_format_instructions(),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
# Map answer_index to ontology
|
||
|
|
# answer_index: 0 -> select None
|
||
|
|
# answer_index: 1 to num_ontologies -> select ontology at (answer_index - 1)
|
||
|
|
state.status = Status.SUCCESS
|
||
|
|
if selector.answer_index == 0:
|
||
|
|
# None selected
|
||
|
|
logger.debug("LLM selected: None (no suitable ontology)")
|
||
|
|
state.current_ontology = NULL_ONTOLOGY
|
||
|
|
elif 1 <= selector.answer_index <= num_ontologies:
|
||
|
|
# Select ontology at index (answer_index - 1) since list is 0-based
|
||
|
|
selected_ontology = ontologies[selector.answer_index - 1]
|
||
|
|
logger.debug(
|
||
|
|
f"LLM selected ontology at index {selector.answer_index}: "
|
||
|
|
f"{selected_ontology.ontology_id} ({selected_ontology.iri})"
|
||
|
|
)
|
||
|
|
state.current_ontology = selected_ontology
|
||
|
|
state.status = Status.SUCCESS
|
||
|
|
else:
|
||
|
|
# This should not happen due to Pydantic validation, but handle gracefully
|
||
|
|
logger.warning(
|
||
|
|
f"Invalid answer_index {selector.answer_index} defaulting to NULL_ONTOLOGY"
|
||
|
|
)
|
||
|
|
state.current_ontology = NULL_ONTOLOGY
|
||
|
|
else:
|
||
|
|
state.current_ontology = NULL_ONTOLOGY
|
||
|
|
|
||
|
|
# Set the initial version if not already set (tracks original version when ontology was selected)
|
||
|
|
if state.current_ontology.initial_version is None:
|
||
|
|
state.current_ontology.initial_version = state.current_ontology.version
|
||
|
|
logger.debug(
|
||
|
|
f"Set initial version for ontology {state.current_ontology.ontology_id}: {state.current_ontology.initial_version}"
|
||
|
|
)
|
||
|
|
|
||
|
|
logger.debug(f"Current ontology set to: {state.current_ontology.ontology_id}")
|
||
|
|
|
||
|
|
return state
|