AI/참고/ontocast-main/ontocast/onto/state.py

import os
from collections import defaultdict
from typing import Any

from pydantic import ConfigDict, Field
from rdflib import URIRef

from ontocast.onto.constants import CHUNK_NULL_IRI, DEFAULT_DOMAIN, ONTOLOGY_NULL_IRI
from ontocast.onto.content_unit import ContentUnit
from ontocast.onto.context import AgentContext, AgentType, ContextManager
from ontocast.onto.enum import FailureStage, RenderMode, Status, WorkflowNode
from ontocast.onto.model import BasePydanticModel, Suggestions
from ontocast.onto.ontology import Ontology
from ontocast.onto.rdfgraph import RDFGraph
from ontocast.onto.sparql_models import GraphUpdate, TripleOp
from ontocast.util import iri2namespace, render_text_hash


class BudgetTracker(BasePydanticModel):
    """Lightweight tracker for LLM usage statistics and generated triples."""

    chars_sent: int = Field(default=0, description="Total characters sent to LLM")
    chars_received: int = Field(
        default=0, description="Total characters received from LLM"
    )
    calls_count: int = Field(default=0, description="Total number of LLM API calls")

    # Triple generation tracking
    ontology_triples_generated: int = Field(
        default=0, description="Total number of triples generated for ontology updates"
    )
    facts_triples_generated: int = Field(
        default=0, description="Total number of triples generated for facts"
    )
    ontology_operations_count: int = Field(
        default=0, description="Total number of ontology update operations"
    )
    facts_operations_count: int = Field(
        default=0, description="Total number of facts update operations"
    )

    def add_usage(self, chars_sent: int, chars_received: int) -> None:
        """Add usage statistics."""
        self.chars_sent += chars_sent
        self.chars_received += chars_received
        self.calls_count += 1

    def add_ontology_update(self, num_operations: int, num_triples: int) -> None:
        """Add ontology update statistics.

        Args:
            num_operations: Number of update operations generated
            num_triples: Number of triples in these operations
        """
        self.ontology_operations_count += num_operations
        self.ontology_triples_generated += num_triples

    def add_facts_update(self, num_operations: int, num_triples: int) -> None:
        """Add facts update statistics.

        Args:
            num_operations: Number of update operations generated
            num_triples: Number of triples in these operations
        """
        self.facts_operations_count += num_operations
        self.facts_triples_generated += num_triples

    def get_summary(self) -> str:
        """Get a summary of LLM usage and generated triples."""
        parts = [
            f"LLM: {self.calls_count} calls, "
            f"{self.chars_sent:,} sent, "
            f"{self.chars_received:,} received",
        ]

        if self.ontology_triples_generated > 0 or self.facts_triples_generated > 0:
            parts.append(
                f"Triples: {self.ontology_triples_generated} ontology, "
                f"{self.facts_triples_generated} facts"
            )

        return " | ".join(parts)


class AgentState(BasePydanticModel):
    """State for the ontology-based knowledge graph agent.

    This class maintains the state of the agent during document processing,
    including input text, content units, ontologies, and workflow status.

    Attributes:
        input_text: Input text to process.
        current_domain: IRI used for forming document namespace.
        doc_hid: An almost unique hash/id for the parent document.
        files: Files to process.
        current_ontology: Current ontology object.
        ontology_addendum: Additional ontology content.
        failure_stage: Stage where failure occurred.
        failure_reason: Reason for failure.
        success_score: Score indicating success level.
        status: Current workflow status.
        node_visits: Number of visits per node.
        max_visits: Maximum number of visits allowed per node.
        max_chunks: Maximum number of source content units to split and process.
    """

    input_text: str = Field(description="Input text", default="")
    current_domain: str = Field(
        description="IRI used for forming document namespace", default=DEFAULT_DOMAIN
    )
    doc_hid: str = Field(
        description="An almost unique hash / id for the parent document of the current unit",
        default="default_doc",
    )
    files: dict[str, bytes] = Field(
        default_factory=lambda: dict(), description="Files to process"
    )
    content_units: list[ContentUnit] = Field(
        default_factory=list,
        description="Pending content units to process.",
    )
    current_content_unit: ContentUnit = Field(
        default_factory=lambda: ContentUnit(
            text="",
            index=0,
            doc_iri=URIRef(CHUNK_NULL_IRI),
        ),
        alias="current_chunk",
        description="Current content unit under processing.",
    )
    current_ontology: Ontology = Field(
        default_factory=lambda: Ontology(
            ontology_id=None,
            title=None,
            description=None,
            graph=RDFGraph(),
            iri=ONTOLOGY_NULL_IRI,
        ),
        description="Ontology object that contain the semantic graph "
        "as well as the description, name, short name, version, "
        "and IRI of the ontology",
    )
    aggregated_facts: RDFGraph = Field(
        description="RDF triples representing aggregated facts "
        "from the current document",
        default_factory=RDFGraph,
    )
    ontology_user_instruction: str = Field(
        description="Specific user instructions for ontology extraction, e.g. `Focus on extracting places`",
        default="",
    )

    facts_user_instruction: str = Field(
        description="Specific user instructions for facts extraction, e.g. `Focus on extracting places`",
        default="",
    )

    dataset: str | None = Field(
        description="Fuseki dataset name for this request (optional)",
        default=None,
    )

    graph_uri_override: str | None = Field(default=None)

    source_url: str | None = Field(
        description="Source URL from JSON input file (for provenance tracking)",
        default=None,
    )

    ontology_updates: list[GraphUpdate] = Field(
        default_factory=list,
        description="A list of graph update that improve the current ontology",
    )

    ontology_updates_applied: list[GraphUpdate] = Field(
        default_factory=list,
        description="A list of graph update that improve the current ontology",
    )

    facts_updates: list[GraphUpdate] = Field(
        default_factory=list,
        description="A list of graph update that improve the current graph of facts (pending)",
    )

    facts_updates_applied: list[GraphUpdate] = Field(
        default_factory=list,
        description="A list of graph update that improve the current graph of facts (applied)",
    )

    parallel_facts_units: list[ContentUnit] = Field(
        default_factory=list,
        description="Successful per-unit facts outputs collected during parallel map phase",
    )

    ontology_units: list[ContentUnit] = Field(
        default_factory=list,
        description="Successful per-unit ontology outputs collected during parallel map phase",
    )
    ontology_provenance_artifact: RDFGraph = Field(
        default_factory=RDFGraph,
        description="Provenance/reification triples stripped from normalized ontology.",
    )

    ontology_addendum: Ontology = Field(
        default_factory=lambda: Ontology(
            ontology_id=None,
            title=None,
            description=None,
            graph=RDFGraph(),
            iri=ONTOLOGY_NULL_IRI,
        ),
        description="Ontology object that contain the semantic graph "
        "as well as the description, name, short name, version, "
        "and IRI of the ontology",
    )
    failure_stage: FailureStage | None = None
    failure_reason: str | None = None

    improvements_suggestions: list[str] = Field(
        description="Itemized concrete and actionable instructions for improvements of extraction of facts/ontology",
        default_factory=list,
    )

    success_score: float = 0.0
    status: Status = Status.SUCCESS
    statuses: dict[WorkflowNode, Status] = Field(
        default_factory=dict, description="Status of each node"
    )
    node_visits: defaultdict[WorkflowNode, int] = Field(
        default_factory=lambda: defaultdict(int),
        description="Number of visits per node",
    )
    max_visits: int = Field(
        default=3, description="Maximum number of visits allowed per node"
    )
    max_chunks: int | None = None
    model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)
    render_mode: RenderMode = Field(
        default=RenderMode.ONTOLOGY_AND_FACTS,
        description=("Rendering mode: ontology, facts, or ontology_and_facts."),
    )
    ontology_max_triples: int | None = Field(
        default=50000,
        description="Maximum number of triples allowed in ontology graph. "
        "Updates that would exceed this limit are skipped with a warning. "
        "Set to None for unlimited.",
    )
    context_manager: ContextManager = Field(
        default_factory=ContextManager,
        description="Context manager for passing information between agents",
    )
    suggestions: Suggestions = Field(
        default_factory=Suggestions,
        description="Context manager for passing information between agents",
    )

    # Budget Tracking
    budget_tracker: BudgetTracker = Field(
        default_factory=BudgetTracker,
        description="Budget statistics tracker (LLM usage and generated triples)",
    )

    def model_post_init(self, __context):
        """Post-initialization hook for the model."""
        pass

    def __init__(self, **kwargs):
        """Initialize the agent state with given keyword arguments."""
        super().__init__(**kwargs)
        self.current_domain = os.getenv("CURRENT_DOMAIN", DEFAULT_DOMAIN)

    def get_node_status(self, node: WorkflowNode) -> Status:
        """Get the status of a workflow node, returning NOT_VISITED if not set."""
        return self.statuses.get(node, Status.NOT_VISITED)

    @property
    def render_ontology(self) -> bool:
        """Whether ontology rendering should run."""
        return self.render_mode in (
            RenderMode.ONTOLOGY,
            RenderMode.ONTOLOGY_AND_FACTS,
        )

    @property
    def render_facts(self) -> bool:
        """Whether facts rendering should run."""
        return self.render_mode in (
            RenderMode.FACTS,
            RenderMode.ONTOLOGY_AND_FACTS,
        )

    def set_node_status(self, node: WorkflowNode, status: Status) -> None:
        """Set the status of a workflow node."""
        self.statuses[node] = status

    def get_content_unit_progress_info(self) -> tuple[int, int]:
        """Get current content unit number and total content units."""
        from ontocast.onto.constants import CHUNK_NULL_IRI

        has_current_content_unit = CHUNK_NULL_IRI not in self.current_content_unit.iri
        current_content_unit_number = 1 if has_current_content_unit else 0
        total_content_units = len(self.content_units)
        return current_content_unit_number, total_content_units

    def get_content_unit_progress_string(self) -> str:
        """Get a formatted string showing content unit progress."""
        current, total = self.get_content_unit_progress_info()
        if total == 0:
            return "no content units"
        return f"content unit {current}/{total}"

    def get_chunk_progress_info(self) -> tuple[int, int]:
        """Backward-compatible wrapper for content unit progress.

        Returns:
            tuple[int, int]: (current_chunk_number, total_chunks)
        """
        return self.get_content_unit_progress_info()

    def get_chunk_progress_string(self) -> str:
        """Backward-compatible wrapper for content unit progress.

        Returns:
            str: Formatted string like "chunk 3/10"
        """
        return self.get_content_unit_progress_string()

    @classmethod
    def render_updated_graph(
        cls, graph: RDFGraph, updates: list[GraphUpdate], max_triples: int | None = None
    ) -> tuple[RDFGraph, bool]:
        """Create a copy of the given graph with all GraphUpdate objects applied.

        This method:
        1. Creates a copy of the input graph
        2. Generates SPARQL queries from all GraphUpdate objects
        3. Executes the queries on the copied graph
        4. Checks if the updated graph exceeds max_triples limit
        5. Returns the updated graph copy, or original if limit exceeded

        Args:
            graph: The RDFGraph to update
            updates: List of GraphUpdate objects to apply
            max_triples: Maximum number of triples allowed. If None, no limit enforced.

        Returns:
            Tuple of (RDFGraph, bool): The updated graph (or original if limit exceeded),
            and a boolean indicating if the update was applied (True) or skipped (False)
        """
        if not updates:
            return graph, True

        # Create a copy of the input graph
        # Use RDFGraph's copy method to preserve type
        updated_graph = RDFGraph()
        for triple in graph:
            updated_graph.add(triple)
        # Copy namespace bindings
        for prefix, namespace in graph.namespaces():
            updated_graph.bind(prefix, namespace)

        all_prefixes = {}
        for graph_update in updates:
            for op in graph_update.triple_operations:
                # Extract prefixes from TripleOp operations
                if isinstance(op, TripleOp) and op.prefixes:
                    all_prefixes.update(op.prefixes)

        # Bind prefixes to the copied graph
        for prefix, uri in all_prefixes.items():
            updated_graph.bind(prefix, uri)

        # Apply each GraphUpdate to the copied graph
        for graph_update in updates:
            # Generate SPARQL queries from the GraphUpdate
            queries = graph_update.generate_sparql_queries()

            # Execute each query on the copied graph
            for query in queries:
                cls._apply_update_query(updated_graph, query)

        # Check if updated graph exceeds max_triples limit
        if max_triples is not None and len(updated_graph) > max_triples:
            import logging

            logger = logging.getLogger(__name__)
            logger.warning(
                f"Ontology update skipped: would exceed limit "
                f"({len(updated_graph)} > {max_triples} triples). "
                f"Original size: {len(graph)} triples."
            )
            return graph, False  # Return original, unchanged

        return updated_graph, True

    @classmethod
    def _apply_update_query(cls, graph: RDFGraph, query: str) -> None:
        """Apply one SPARQL update query, recovering common LLM compound outputs."""
        try:
            graph.update(query)
            return
        except Exception as exc:
            split_queries = cls._split_compound_sparql_updates(query, exc)
            if not split_queries:
                raise
            for split_query in split_queries:
                graph.update(split_query)

    @staticmethod
    def _split_compound_sparql_updates(
        query: str, parse_error: Exception
    ) -> list[str] | None:
        """Split concatenated top-level UPDATE statements if parser rejects input.

        LLM outputs sometimes concatenate multiple update statements (e.g. two
        top-level INSERT blocks) into a single custom SPARQL string without a
        separator accepted by the parser. We only recover in that specific case.
        """
        message = str(parse_error)
        if "Expected end of text, found" not in message:
            return None

        prefixes: list[str] = []
        body_lines: list[str] = []
        for raw_line in query.splitlines():
            stripped = raw_line.strip()
            if stripped.upper().startswith("PREFIX "):
                prefixes.append(stripped)
            elif stripped:
                body_lines.append(stripped)

        if not body_lines:
            return None

        top_level_ops = ("INSERT", "DELETE", "WITH")
        segments: list[list[str]] = []
        current: list[str] = []
        for line in body_lines:
            if current and line.upper().startswith(top_level_ops):
                segments.append(current)
                current = [line]
            else:
                current.append(line)
        if current:
            segments.append(current)

        if len(segments) <= 1:
            return None

        prefix_block = "\n".join(prefixes)
        rebuilt = []
        for segment in segments:
            segment_block = "\n".join(segment)
            rebuilt.append(
                f"{prefix_block}\n{segment_block}" if prefix_block else segment_block
            )
        return rebuilt

    def render_uptodate_ontology(self) -> Ontology:
        """Create a copy of the current ontology with all GraphUpdate objects applied.

        This method:
        1. Creates a copy of the current ontology
        2. Generates SPARQL queries from all GraphUpdate objects
        3. Executes the queries on the copied ontology graph
        4. Checks if the updated graph exceeds max_triples limit
        5. Sets the current hash as parent_hash in the updated ontology
        6. Computes a new hash for the updated ontology
        7. Syncs properties to ensure object fields are updated
        8. Returns the updated ontology copy, or original if limit exceeded

        Returns:
            Ontology: A copy of the current ontology with all updates applied and
            a new hash generated, with the previous hash set as parent.
            Returns original ontology if update would exceed max_triples limit.
        """
        if not self.ontology_updates:
            return self.current_ontology

        # Use the generalized function to update the graph
        updated_graph, was_applied = self.render_updated_graph(
            self.current_ontology.graph,
            self.ontology_updates,
            max_triples=self.ontology_max_triples,
        )

        # If graph wasn't updated (limit exceeded), return original ontology
        if not was_applied:
            return self.current_ontology

        return self.current_ontology.derive_updated_version(updated_graph)

    def update_ontology(self) -> None:
        """Update the current ontology with all GraphUpdate objects and clear the updates list.

        This method:
        1. Uses render_uptodate_ontology() to get an updated copy
        2. Replaces the current ontology with the updated copy
        3. Clears the ontology_updates list

        Note: Version update is deferred to aggregate_serialize() to update only once at the end.
        """
        if not self.ontology_updates:
            return

        # Get the updated ontology copy
        updated_ontology = self.render_uptodate_ontology()

        # Replace the current ontology with the updated copy
        self.current_ontology = updated_ontology

        # Clear the updates list
        self.ontology_updates_applied += self.ontology_updates
        self.ontology_updates = []

    def render_uptodate_facts(self) -> RDFGraph:
        """Create a copy of the current content unit graph with facts updates applied.

        This method:
        1. Creates a copy of the current content unit's graph
        2. Generates SPARQL queries from all facts GraphUpdate objects
        3. Executes the queries on the copied graph
        4. Returns the updated graph copy

        Returns:
            RDFGraph: A copy of the current chunk's graph with all facts updates applied
        """
        if not self.facts_updates:
            return self.current_content_unit.graph

        # Use the generalized function to update the graph
        updated_graph, _ = self.render_updated_graph(
            self.current_content_unit.graph, self.facts_updates, max_triples=None
        )
        return updated_graph

    def update_facts(self) -> None:
        """Update current content unit graph with facts updates and clear the updates list.

        This method:
        1. Uses render_uptodate_facts() to get an updated copy
        2. Replaces the current content unit graph with the updated copy
        3. Clears the facts_updates list
        """
        if not self.facts_updates:
            return

        # Get the updated graph copy
        updated_graph = self.render_uptodate_facts()

        # Replace the current chunk's graph with the updated copy
        self.current_content_unit.graph = updated_graph

        # Clear the updates list
        self.facts_updates_applied += self.facts_updates
        self.facts_updates = []

    def generate_ontology_updates_markdown(self) -> str:
        """Generate a markdown string representing the chain of ontology updates.

        Returns:
            Markdown-formatted string showing all pending ontology updates.
            Returns empty string if no updates are pending.
        """
        if not self.ontology_updates:
            return ""

        markdown_parts = []
        for i, graph_update in enumerate(self.ontology_updates, 1):
            diff_summary = graph_update.generate_diff_summary()
            if diff_summary:
                markdown_parts.append(f"## Update {i}")
                markdown_parts.append(diff_summary)

            markdown_parts.append("")

            # Add separator between updates (except for the last one)
            if i < len(self.ontology_updates):
                markdown_parts.append("---")
                markdown_parts.append("")

        return "\n".join(markdown_parts)

    def set_text(self, text):
        """Set the input text and generate document hash.

        Args:
            text: The input text to set.
        """
        self.input_text = text
        self.doc_hid = render_text_hash(self.input_text)

    def set_failure(self, stage: FailureStage, reason: str, success_score: float = 0.0):
        """Set failure state with stage and reason.

        Args:
            stage: The stage where the failure occurred.
            reason: The reason for the failure.
            success_score: The success score at failure (default: 0.0).
        """
        self.failure_stage = stage
        self.failure_reason = reason
        self.success_score = success_score
        self.status = Status.FAILED

    def clear_failure(self):
        """Clear failure state and set status to success."""
        self.failure_stage = None
        self.failure_reason = None
        self.success_score = 0.0
        self.status = Status.SUCCESS

    @property
    def doc_iri(self) -> URIRef:
        """Get the document IRI.

        Returns:
            str: The document IRI.
        """
        return URIRef(f"{self.current_domain}/doc/{self.doc_hid}")

    @property
    def doc_namespace(self):
        """Get the document namespace.

        Returns:
            str: The document namespace.
        """
        return iri2namespace(self.doc_iri, ontology=False)

    @property
    def graph_uri(self):
        if self.graph_uri_override is not None:
            return self.graph_uri_override
        return self.doc_namespace

    @property
    def ontology_id(self):
        """Get the document namespace.

        Returns:
            str: The document namespace.
        """
        return self.current_ontology.ontology_id

    def get_context_for_agent(self, agent_type: AgentType) -> AgentContext:
        """Get or create context for a specific agent.

        Args:
            agent_type: Type of agent (renderer, critic, etc.).

        Returns:
            AgentContext: The context for the agent.
        """
        existing_context = self.context_manager.get_latest_context_by_agent(agent_type)

        if existing_context:
            return existing_context

        # Create new context if none exists
        return self.context_manager.create_context(agent_type=agent_type)

    def update_context_for_agent(
        self,
        agent_type: AgentType,
        ontology_version: Any | None = None,
        facts_version: Any | None = None,
        ontology_operations: list[Any] | None = None,
        facts_operations: list[Any] | None = None,
        ontology_critique: dict[str, Any] | None = None,
        facts_critique: dict[str, Any] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> AgentContext:
        """Update context for a specific agent.

        Args:
            agent_type: Name of the agent updating context.
            ontology_version: New ontology version if available.
            facts_version: New facts version if available.
            ontology_operations: New ontology operations if available.
            facts_operations: New facts operations if available.
            ontology_critique: New ontology critique if available.
            facts_critique: New facts critique if available.
            metadata: Additional metadata for the context.

        Returns:
            AgentContext: The updated context.
        """
        return self.context_manager.update_context(
            agent_type=agent_type,
            ontology_version=ontology_version,
            facts_version=facts_version,
            ontology_operations=ontology_operations,
            facts_operations=facts_operations,
            ontology_critique=ontology_critique,
            facts_critique=facts_critique,
            metadata=metadata,
        )

    def get_context_summary_for_agent(self, agent_type: AgentType) -> str:
        """Get a context summary for a specific agent.

        Args:
            agent_type: Name of the agent requesting context summary.

        Returns:
            str: A formatted context summary.
        """
        context = self.context_manager.get_latest_context_by_agent(agent_type)
        if not context:
            return "No context available for this agent."

        return context.get_full_context_summary()