mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
Properly saving events with session id to local sqlite db
This commit is contained in:
@@ -1,8 +1,10 @@
|
||||
import warnings
|
||||
|
||||
from crewai.agent import Agent
|
||||
from crewai.crew import Crew
|
||||
from crewai.flow.flow import Flow
|
||||
from crewai.llm import LLM
|
||||
from crewai.logging.event_logger import EventLogger
|
||||
from crewai.pipeline import Pipeline
|
||||
from crewai.process import Process
|
||||
from crewai.routers import Router
|
||||
|
||||
@@ -215,7 +215,6 @@ class BaseAgent(ABC, BaseModel):
|
||||
def serialize(self) -> Dict[str, Any]:
|
||||
"""Serialize the BaseAgent into a dictionary excluding complex objects."""
|
||||
|
||||
# Define attributes to exclude from serialization
|
||||
exclude = {
|
||||
"_logger",
|
||||
"_rpm_controller",
|
||||
@@ -229,8 +228,8 @@ class BaseAgent(ABC, BaseModel):
|
||||
# Add any other complex attributes that should be excluded
|
||||
}
|
||||
|
||||
# Use model_dump or similar to get a dictionary representation
|
||||
serialized_data = self.model_dump(exclude=exclude)
|
||||
serialized_data = {k: v for k, v in serialized_data.items() if v is not None}
|
||||
|
||||
# Add any additional serialization logic if needed
|
||||
serialized_data["role"] = self.role
|
||||
@@ -240,6 +239,9 @@ class BaseAgent(ABC, BaseModel):
|
||||
[str(tool) for tool in self.tools] if self.tools else None
|
||||
)
|
||||
|
||||
# Include the UUID as a string
|
||||
serialized_data["id"] = str(self.id)
|
||||
|
||||
return serialized_data
|
||||
|
||||
def copy(self: T) -> T: # type: ignore # Signature of "copy" incompatible with supertype "BaseModel"
|
||||
|
||||
@@ -922,6 +922,9 @@ class Crew(BaseModel):
|
||||
serialized_data["agents"] = serialized_agents
|
||||
serialized_data["tasks"] = serialized_tasks
|
||||
|
||||
# Include the UUID as a string
|
||||
serialized_data["id"] = str(self.id)
|
||||
|
||||
return serialized_data
|
||||
|
||||
# TODO: Come back and use the new _serialize method
|
||||
|
||||
@@ -1,23 +1,39 @@
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from typing import Any, Optional
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from sqlalchemy import JSON, Column, DateTime, Integer, String, create_engine
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
from sqlalchemy.orm import Session as SQLAlchemySession
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from crewai.utilities.event_emitter import crew_events
|
||||
|
||||
load_dotenv()
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class Session:
|
||||
_session_id: Optional[str] = None
|
||||
|
||||
@classmethod
|
||||
def get_session_id(cls) -> str:
|
||||
if cls._session_id is None:
|
||||
cls._session_id = str(uuid.uuid4()) # Generate a new UUID
|
||||
print(f"Generated new session ID: {cls._session_id}")
|
||||
return cls._session_id
|
||||
|
||||
|
||||
class Event(Base):
|
||||
__tablename__ = "events"
|
||||
id = Column(Integer, primary_key=True)
|
||||
event_name = Column(String)
|
||||
timestamp = Column(DateTime, default=datetime.now(timezone.utc))
|
||||
crew_id = Column(String)
|
||||
session_id = Column(String)
|
||||
data = Column(JSON)
|
||||
error_type = Column(String)
|
||||
error_message = Column(String)
|
||||
@@ -30,24 +46,49 @@ Base.metadata.create_all(engine)
|
||||
SessionLocal = sessionmaker(bind=engine) # Use a different name to avoid confusion
|
||||
|
||||
# Create a session instance
|
||||
session: Session = SessionLocal()
|
||||
session: SQLAlchemySession = SessionLocal()
|
||||
|
||||
|
||||
class EventLogger:
|
||||
def __init__(self, session: Session):
|
||||
def __init__(self, session: SQLAlchemySession):
|
||||
self.session = session
|
||||
|
||||
def log_event(self, event_name: str, *args: Any, **kwargs: Any) -> None:
|
||||
def log_event(self, *args: Any, **kwargs: Any) -> None:
|
||||
# Extract event name from kwargs
|
||||
event_name = kwargs.pop("event", "unknown_event")
|
||||
print("Logging event:", event_name)
|
||||
print("Args:", args)
|
||||
print("Kwargs:", kwargs)
|
||||
|
||||
# Check if args is a single dictionary and unpack it
|
||||
if len(args) == 1 and isinstance(args[0], dict):
|
||||
args_dict = args[0]
|
||||
else:
|
||||
# Convert args to a dictionary with keys like 'arg0', 'arg1', etc.
|
||||
args_dict = {f"arg{i}": arg for i, arg in enumerate(args)}
|
||||
|
||||
# Merge args_dict and kwargs into a single dictionary
|
||||
data = {**args_dict, **kwargs}
|
||||
|
||||
print("Data:", data)
|
||||
|
||||
event = Event(
|
||||
event_name=event_name,
|
||||
crew_id=kwargs.get("crew_id", ""),
|
||||
data=kwargs,
|
||||
session_id=Session.get_session_id(),
|
||||
data=json.dumps(data),
|
||||
error_type=kwargs.get("error_type", ""),
|
||||
error_message=kwargs.get("error_message", ""),
|
||||
traceback=kwargs.get("traceback", ""),
|
||||
)
|
||||
|
||||
self.session.add(event)
|
||||
self.session.commit()
|
||||
print("Successfully logged event:", event_name)
|
||||
|
||||
|
||||
event_logger = EventLogger(session)
|
||||
|
||||
|
||||
print("Connecting event_logger to all signals")
|
||||
crew_events.on("*", event_logger.log_event)
|
||||
print("Connected event_logger to all signals")
|
||||
|
||||
@@ -20,21 +20,32 @@ class CrewEventEmitter:
|
||||
def __init__(self):
|
||||
self._all_signal = signal("all")
|
||||
|
||||
def on(self, event_name: CrewEvents, callback: Callable) -> None:
|
||||
if event_name == "*":
|
||||
self._all_signal.connect(callback)
|
||||
def on(self, event_name: CrewEvents | str, callback: Callable) -> None:
|
||||
print("Connecting signal:", event_name)
|
||||
if event_name == "*" or event_name == "all":
|
||||
self._all_signal.connect(callback, weak=False)
|
||||
print("Connected to all_signal")
|
||||
else:
|
||||
signal(event_name.value).connect(callback)
|
||||
signal(
|
||||
event_name.value if isinstance(event_name, CrewEvents) else event_name
|
||||
).connect(callback, weak=False)
|
||||
|
||||
def emit(self, event_name: CrewEvents, *args: Any, **kwargs: Any) -> None:
|
||||
print(f"Emitting signal: {event_name.value}")
|
||||
print("args", args)
|
||||
print("kwargs", kwargs)
|
||||
signal(event_name.value).send(*args, **kwargs)
|
||||
self._all_signal.send(event_name, *args, **kwargs)
|
||||
print(f"Emitting all signal for: {event_name.value}")
|
||||
self._all_signal.send(*args, event=event_name.value, **kwargs)
|
||||
|
||||
|
||||
crew_events = CrewEventEmitter()
|
||||
|
||||
|
||||
def emit(event_name: CrewEvents, *args: Any, **kwargs: Any) -> None:
|
||||
print("Calling emit", event_name)
|
||||
print("Args:", args)
|
||||
print("Kwargs:", kwargs)
|
||||
try:
|
||||
crew_events.emit(event_name, *args, **kwargs)
|
||||
except Exception as e:
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
# event_helpers.py
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
|
||||
from crewai.crew import Crew
|
||||
from crewai.utilities.event_emitter import CrewEvents, emit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from crewai.crew import Crew
|
||||
|
||||
|
||||
def emit_crew_start(
|
||||
crew: Crew,
|
||||
crew: "Crew", # Use a forward reference
|
||||
inputs: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
serialized_crew = crew.serialize()
|
||||
|
||||
Reference in New Issue
Block a user