""" Database models for the agent monitoring system. """ from datetime import datetime, timezone import json from sqlalchemy import Column, Integer, String, Float, Boolean, DateTime, Text, ForeignKey, Table, UniqueConstraint, Index from sqlalchemy.orm import relationship from sqlalchemy.types import JSON, TypeDecorator from sqlalchemy.ext.declarative import declarative_base import uuid Base = declarative_base() class SafeJSON(TypeDecorator): """Custom JSON type that handles circular references using default=str""" impl = Text def process_bind_param(self, value, dialect): if value is not None: return json.dumps(value, default=str) return value def process_result_value(self, value, dialect): if value is not None: return json.loads(value) return value class Trace(Base): """Model for storing agent traces (conversations, interactions, etc.).""" __tablename__ = "traces" id = Column(Integer, primary_key=True, index=True) trace_id = Column(String(36), unique=True, index=True, default=lambda: str(uuid.uuid4())) filename = Column(String(255), nullable=True, index=True) title = Column(String(255), nullable=True) description = Column(Text, nullable=True) content = Column(Text, nullable=True) # Full trace content content_hash = Column(String(64), nullable=True, index=True) # Hash of content for deduplication upload_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc)) update_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) uploader = Column(String(255), nullable=True) trace_type = Column(String(50), nullable=True) # e.g., 'conversation', 'code_execution', etc. trace_source = Column(String(50), nullable=True) # e.g., 'user_upload', 'api', 'generated' character_count = Column(Integer, default=0) turn_count = Column(Integer, default=0) status = Column(String(50), default="uploaded") # uploaded, processed, analyzed, etc. processing_method = Column(String(50), nullable=True) # e.g., 'sliding_window', 'single_pass', etc. tags = Column(JSON, nullable=True) # Store tags as JSON array trace_metadata = Column(JSON, nullable=True) # Additional metadata as JSON # Relationships knowledge_graphs = relationship("KnowledgeGraph", back_populates="trace", foreign_keys="KnowledgeGraph.trace_id", cascade="all, delete-orphan") __table_args__ = ( UniqueConstraint('trace_id', name='uix_trace_id'), Index('idx_trace_content_hash', 'content_hash'), Index('idx_trace_title', 'title'), Index('idx_trace_status', 'status'), ) def to_dict(self): """Convert to dictionary representation.""" return { "id": self.id, "trace_id": self.trace_id, "filename": self.filename, "title": self.title, "description": self.description, "upload_timestamp": self.upload_timestamp.isoformat() if self.upload_timestamp else None, "update_timestamp": self.update_timestamp.isoformat() if self.update_timestamp else None, "uploader": self.uploader, "trace_type": self.trace_type, "trace_source": self.trace_source, "character_count": self.character_count, "turn_count": self.turn_count, "status": self.status, "processing_method": self.processing_method, "tags": self.tags, "metadata": self.trace_metadata, "knowledge_graph_count": len(self.knowledge_graphs) if self.knowledge_graphs else 0 } @classmethod def from_content(cls, content, filename=None, title=None, description=None, trace_type=None, trace_source="user_upload", uploader=None, tags=None, trace_metadata=None): """Create a Trace instance from content.""" import hashlib trace = cls() trace.trace_id = str(uuid.uuid4()) trace.filename = filename trace.title = title or f"Trace {trace.trace_id[:8]}" trace.description = description trace.content = content # Calculate content hash for deduplication if content: content_hash = hashlib.sha256(content.encode('utf-8')).hexdigest() trace.content_hash = content_hash # Set character count trace.character_count = len(content) # Estimate turn count (approximate) turn_markers = [ "user:", "assistant:", "system:", "human:", "ai:", "User:", "Assistant:", "System:", "Human:", "AI:" ] turn_count = 0 for marker in turn_markers: turn_count += content.count(marker) trace.turn_count = max(1, turn_count) # At least 1 turn trace.trace_type = trace_type trace.trace_source = trace_source trace.uploader = uploader trace.tags = tags or [] trace.trace_metadata = trace_metadata or {} trace.status = "uploaded" return trace class KnowledgeGraph(Base): """Model for storing knowledge graphs.""" __tablename__ = "knowledge_graphs" id = Column(Integer, primary_key=True, index=True) filename = Column(String(255), unique=True, index=True) creation_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc)) update_timestamp = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) entity_count = Column(Integer, default=0) relation_count = Column(Integer, default=0) _graph_data = Column("graph_data", Text, nullable=True) # Underlying TEXT field status = Column(String(50), default="created", nullable=False) # Status of processing: created, enriched, perturbed, causal # Add fields for trace and window tracking trace_id = Column(String(36), ForeignKey("traces.trace_id"), nullable=True, index=True, comment="ID to group knowledge graphs from the same trace") window_index = Column(Integer, nullable=True, comment="Sequential index of window within a trace") window_total = Column(Integer, nullable=True, comment="Total number of windows in the trace") window_start_char = Column(Integer, nullable=True, comment="Starting character position in the original trace") window_end_char = Column(Integer, nullable=True, comment="Ending character position in the original trace") processing_run_id = Column(String(36), nullable=True, index=True, comment="ID to distinguish multiple processing runs of the same trace") # Relationships entities = relationship("Entity", back_populates="graph", cascade="all, delete-orphan") relations = relationship("Relation", back_populates="graph", cascade="all, delete-orphan") trace = relationship("Trace", back_populates="knowledge_graphs", foreign_keys=[trace_id]) prompt_reconstructions = relationship( "PromptReconstruction", back_populates="knowledge_graph", cascade="all, delete-orphan" ) perturbation_tests = relationship("PerturbationTest", back_populates="knowledge_graph", cascade="all, delete-orphan") causal_analyses = relationship("CausalAnalysis", back_populates="knowledge_graph", cascade="all, delete-orphan") __table_args__ = ( UniqueConstraint('filename', name='uix_knowledge_graph_filename'), ) @property def graph_data(self): """Get the graph_data as a parsed JSON object""" if self._graph_data is None: return None if isinstance(self._graph_data, dict): # Already a dictionary, return as is return self._graph_data # Try to parse as JSON try: return json.loads(self._graph_data) except: # If parsing fails, return None return None @graph_data.setter def graph_data(self, value): """Set graph_data, converting to a JSON string if it's a dictionary""" if value is None: self._graph_data = None elif isinstance(value, dict): self._graph_data = json.dumps(value) else: # Assume it's already a string self._graph_data = value @property def graph_content(self): """Get the graph content from graph_data field""" # Return graph_data return self.graph_data or {} @graph_content.setter def graph_content(self, data): """Set graph content from a dictionary.""" self.graph_data = data # Update counts if isinstance(data, dict): if 'entities' in data and isinstance(data['entities'], list): self.entity_count = len(data['entities']) if 'relations' in data and isinstance(data['relations'], list): self.relation_count = len(data['relations']) def get_entities_from_content(self): """Get entities directly from content field.""" data = self.graph_content entities = data.get('entities', []) if isinstance(data, dict) else [] return entities def get_relations_from_content(self): """Get relations directly from content field.""" data = self.graph_content relations = data.get('relations', []) if isinstance(data, dict) else [] return relations def get_all_entities(self, session=None): """ Get all entities, preferring database entities if available. If no database entities exist, falls back to content entities. If session is provided, queries database entities, otherwise returns content entities. """ if session: db_entities = session.query(Entity).filter_by(graph_id=self.id).all() if db_entities: return [entity.to_dict() for entity in db_entities] return self.get_entities_from_content() def get_all_relations(self, session=None): """ Get all relations, preferring database relations if available. If no database relations exist, falls back to content relations. If session is provided, queries database relations, otherwise returns content relations. """ if session: db_relations = session.query(Relation).filter_by(graph_id=self.id).all() if db_relations: return [relation.to_dict() for relation in db_relations] return self.get_relations_from_content() def to_dict(self): """Convert to dictionary representation.""" result = { "id": self.id, "filename": self.filename, "creation_timestamp": self.creation_timestamp.isoformat(), "entity_count": self.entity_count, "relation_count": self.relation_count, } return result @classmethod def from_dict(cls, data): """Create a KnowledgeGraph instance from a dictionary representation.""" kg = cls() kg.filename = data.get('filename') # Store content as JSON kg.content = json.dumps(data) return kg class Entity(Base): """Model for storing knowledge graph entities.""" __tablename__ = "entities" id = Column(Integer, primary_key=True, index=True) graph_id = Column(Integer, ForeignKey("knowledge_graphs.id")) entity_id = Column(String(255), index=True) # Original entity ID in the graph type = Column(String(255)) name = Column(String(255)) properties = Column(JSON) # Relationships graph = relationship("KnowledgeGraph", back_populates="entities") source_relations = relationship("Relation", foreign_keys="Relation.source_id", back_populates="source") target_relations = relationship("Relation", foreign_keys="Relation.target_id", back_populates="target") # Add a composite unique constraint to ensure entity_id is unique per graph __table_args__ = ( UniqueConstraint('graph_id', 'entity_id', name='uix_entity_graph_id_entity_id'), ) def to_dict(self): """Convert to dictionary representation.""" result = { "id": self.entity_id, "type": self.type, "name": self.name, "properties": self.properties or {} } return result @classmethod def from_dict(cls, data, graph_id): """Create an Entity instance from a dictionary.""" entity = cls() entity.graph_id = graph_id entity.entity_id = data.get('id') entity.type = data.get('type') entity.name = data.get('name') entity.properties = data.get('properties') return entity class Relation(Base): """Model for storing knowledge graph relations.""" __tablename__ = "relations" id = Column(Integer, primary_key=True, index=True) graph_id = Column(Integer, ForeignKey("knowledge_graphs.id")) relation_id = Column(String(255), index=True) # Original relation ID in the graph type = Column(String(255)) source_id = Column(Integer, ForeignKey("entities.id")) target_id = Column(Integer, ForeignKey("entities.id")) properties = Column(JSON) # Relationships graph = relationship("KnowledgeGraph", back_populates="relations") source = relationship("Entity", foreign_keys=[source_id], back_populates="source_relations") target = relationship("Entity", foreign_keys=[target_id], back_populates="target_relations") # Add a composite unique constraint to ensure relation_id is unique per graph __table_args__ = ( UniqueConstraint('graph_id', 'relation_id', name='uix_relation_graph_id_relation_id'), ) def to_dict(self): """Convert to dictionary representation.""" result = { "id": self.relation_id, "type": self.type, "source": self.source.entity_id if self.source else None, "target": self.target.entity_id if self.target else None, "properties": self.properties or {} } return result @classmethod def from_dict(cls, data, graph_id, source_entity=None, target_entity=None): """Create a Relation instance from a dictionary.""" relation = cls() relation.graph_id = graph_id relation.relation_id = data.get('id') relation.type = data.get('type') # Set source and target if source_entity: relation.source_id = source_entity.id if target_entity: relation.target_id = target_entity.id # Set properties relation.properties = data.get('properties') return relation class PromptReconstruction(Base): """Model for storing prompt reconstruction results.""" __tablename__ = "prompt_reconstructions" id = Column(Integer, primary_key=True) knowledge_graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"), nullable=False) relation_id = Column(String(255), nullable=False) reconstructed_prompt = Column(Text) dependencies = Column(JSON) created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) # Relationships knowledge_graph = relationship("KnowledgeGraph", back_populates="prompt_reconstructions") perturbation_tests = relationship("PerturbationTest", back_populates="prompt_reconstruction") def to_dict(self): return { "id": self.id, "knowledge_graph_id": self.knowledge_graph_id, "relation_id": self.relation_id, "reconstructed_prompt": self.reconstructed_prompt, "dependencies": self.dependencies, "created_at": self.created_at.isoformat() if self.created_at else None, "updated_at": self.updated_at.isoformat() if self.updated_at else None } class PerturbationTest(Base): """Model for storing perturbation test results.""" __tablename__ = "perturbation_tests" id = Column(Integer, primary_key=True) knowledge_graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"), nullable=False) prompt_reconstruction_id = Column(Integer, ForeignKey("prompt_reconstructions.id"), nullable=False) relation_id = Column(String(255), nullable=False) perturbation_type = Column(String(50), nullable=False) # e.g., 'entity_removal', 'relation_removal' perturbation_set_id = Column(String(64), nullable=False, index=True) test_result = Column(JSON) perturbation_score = Column(Float) test_metadata = Column(JSON) created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) # Relationships knowledge_graph = relationship("KnowledgeGraph", back_populates="perturbation_tests") prompt_reconstruction = relationship("PromptReconstruction", back_populates="perturbation_tests") def to_dict(self): return { "id": self.id, "knowledge_graph_id": self.knowledge_graph_id, "prompt_reconstruction_id": self.prompt_reconstruction_id, "relation_id": self.relation_id, "perturbation_type": self.perturbation_type, "perturbation_set_id": self.perturbation_set_id, "test_result": self.test_result, "perturbation_score": self.perturbation_score, "test_metadata": self.test_metadata, "created_at": self.created_at.isoformat() if self.created_at else None, "updated_at": self.updated_at.isoformat() if self.updated_at else None } class CausalAnalysis(Base): """Model for storing causal analysis results.""" __tablename__ = "causal_analyses" id = Column(Integer, primary_key=True) knowledge_graph_id = Column(Integer, ForeignKey("knowledge_graphs.id"), nullable=False) perturbation_set_id = Column(String(64), nullable=False, index=True) # Analysis method and results analysis_method = Column(String(50), nullable=False) # e.g., 'graph', 'component', 'dowhy' analysis_result = Column(JSON) # Store the full analysis result causal_score = Column(Float) # Store the numerical causal score analysis_metadata = Column(JSON) # Store additional metadata about the analysis # Timestamps created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) # Relationships knowledge_graph = relationship("KnowledgeGraph", back_populates="causal_analyses") # Indexes __table_args__ = ( Index("idx_causal_analyses_kgid", "knowledge_graph_id"), Index("idx_causal_analyses_method", "analysis_method"), Index("idx_causal_analyses_setid", "perturbation_set_id"), ) def to_dict(self): return { "id": self.id, "knowledge_graph_id": self.knowledge_graph_id, "perturbation_set_id": self.perturbation_set_id, "analysis_method": self.analysis_method, "analysis_result": self.analysis_result, "causal_score": self.causal_score, "analysis_metadata": self.analysis_metadata, "created_at": self.created_at.isoformat() if self.created_at else None, "updated_at": self.updated_at.isoformat() if self.updated_at else None } class ObservabilityConnection(Base): """Model for storing AI observability platform connections.""" __tablename__ = "observability_connections" id = Column(Integer, primary_key=True, index=True) connection_id = Column(String(36), unique=True, index=True, default=lambda: str(uuid.uuid4())) platform = Column(String(50), nullable=False) # langfuse, langsmith, etc. public_key = Column(Text, nullable=False) # Encrypted API key secret_key = Column(Text, nullable=True) # Encrypted secret key (for Langfuse) host = Column(String(255), nullable=True) # Host URL projects = Column(JSON, nullable=True) # Available projects from the platform status = Column(String(50), default="connected") connected_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) last_sync = Column(DateTime, nullable=True) created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), onupdate=lambda: datetime.now(timezone.utc)) # Relationships fetched_traces = relationship("FetchedTrace", back_populates="connection", cascade="all, delete-orphan") def to_dict(self): return { "id": self.connection_id, "platform": self.platform, "status": self.status, "connected_at": self.connected_at.isoformat() if self.connected_at else None, "last_sync": self.last_sync.isoformat() if self.last_sync else None, "host": self.host, "projects": self.projects or [] } class FetchedTrace(Base): """Model for storing fetched traces from observability platforms.""" __tablename__ = "fetched_traces" id = Column(Integer, primary_key=True, index=True) trace_id = Column(String(255), nullable=False, index=True) # Original trace ID from platform name = Column(String(255), nullable=False) platform = Column(String(50), nullable=False) connection_id = Column(String(36), ForeignKey("observability_connections.connection_id"), nullable=False) project_name = Column(String(255), nullable=True, index=True) # Project name for LangSmith, null for Langfuse data = Column(SafeJSON, nullable=True) # Full trace data fetched_at = Column(DateTime, default=lambda: datetime.now(timezone.utc)) imported = Column(Boolean, default=False) imported_at = Column(DateTime, nullable=True) imported_trace_id = Column(String(36), nullable=True) # Reference to imported trace # Relationships connection = relationship("ObservabilityConnection", back_populates="fetched_traces") __table_args__ = ( UniqueConstraint('trace_id', 'connection_id', name='uix_fetched_trace_id_connection'), ) def _extract_generated_timestamp(self): """Extract the actual generated timestamp from trace data based on platform.""" if not self.data: return None if self.platform == "langfuse": # For Langfuse, find the earliest timestamp from traces traces = self.data.get("traces", []) if traces: timestamps = [] for trace in traces: if isinstance(trace, dict): # Check for various timestamp fields in Langfuse traces for ts_field in ["timestamp", "startTime", "createdAt"]: if ts_field in trace: timestamps.append(trace[ts_field]) break if timestamps: return min(timestamps) # Fallback to session info or other timestamps session_info = self.data.get("session_info", {}) if session_info and "createdAt" in session_info: return session_info["createdAt"] # Other fallback fields at top level for field in ["timestamp", "createdAt", "startTime"]: if field in self.data: return self.data[field] elif self.platform == "langsmith": # For LangSmith, find the earliest start_time from traces traces = self.data.get("traces", []) if traces: start_times = [] for trace in traces: if isinstance(trace, dict) and "start_time" in trace: start_times.append(trace["start_time"]) if start_times: return min(start_times) # Fallback to other timestamp fields for field in ["timestamp", "start_time", "created_at"]: if field in self.data: return self.data[field] return None def to_dict(self, preview=True): data = self.data original_stats = {} if data: # Calculate original data statistics import json original_json_str = json.dumps(data, ensure_ascii=False) original_stats = { "original_character_count": len(original_json_str), "original_line_count": original_json_str.count('\n') + 1, "original_size_kb": round(len(original_json_str) / 1024, 2) } if preview: # Truncate long strings to prevent browser crashes but preserve full structure from backend.routers.observability import truncate_long_strings data = truncate_long_strings(data, max_string_length=500) # Extract generated timestamp generated_timestamp = self._extract_generated_timestamp() result = { "id": self.trace_id, "name": self.name, "platform": self.platform, "fetched_at": self.fetched_at.isoformat() if self.fetched_at else None, "generated_timestamp": generated_timestamp, "imported": self.imported, "imported_at": self.imported_at.isoformat() if self.imported_at else None, "data": data } # Add original statistics to the result result.update(original_stats) return result def get_full_data(self): """Get full original data for download (no limitations)""" return { "id": self.trace_id, "name": self.name, "platform": self.platform, "fetched_at": self.fetched_at.isoformat() if self.fetched_at else None, "imported": self.imported, "imported_at": self.imported_at.isoformat() if self.imported_at else None, "data": self.data # Full original data }