- Implementation of specialist execution api, including SSE protocol
- eveai_chat becomes deprecated and should be replaced with SSE - Adaptation of STANDARD_RAG specialist - Base class definition allowing to realise specialists with crewai framework - Implementation of SPIN_SPECIALIST - Implementation of test app for testing specialists (test_specialist_client). Also serves as an example for future SSE-based client - Improvements to startup scripts to better handle and scale multiple connections - Small improvements to the interaction forms and views - Caching implementation improved and augmented with additional caches
This commit is contained in:
Josako
296
eveai_chat_workers/specialists/SPIN_SPECIALIST/1_0.py
Normal file
296
eveai_chat_workers/specialists/SPIN_SPECIALIST/1_0.py
Normal file
@@ -0,0 +1,296 @@
|
||||
import json
|
||||
from os import wait
|
||||
from typing import Optional, List
|
||||
|
||||
from crewai.flow.flow import start, listen, and_
|
||||
from flask import current_app
|
||||
from gevent import sleep
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from common.extensions import cache_manager
|
||||
from common.models.user import Tenant
|
||||
from common.utils.business_event_context import current_event
|
||||
from eveai_chat_workers.retrievers.retriever_typing import RetrieverArguments
|
||||
from eveai_chat_workers.specialists.crewai_base_specialist import CrewAIBaseSpecialistExecutor
|
||||
from eveai_chat_workers.specialists.specialist_typing import SpecialistResult, SpecialistArguments
|
||||
from eveai_chat_workers.outputs.identification.identification_v1_0 import LeadInfoOutput
|
||||
from eveai_chat_workers.outputs.spin.spin_v1_0 import SPINOutput
|
||||
from eveai_chat_workers.outputs.rag.rag_v1_0 import RAGOutput
|
||||
from eveai_chat_workers.specialists.crewai_base_classes import EveAICrewAICrew, EveAICrewAIFlow, EveAIFlowState
|
||||
from common.utils.pydantic_utils import flatten_pydantic_model
|
||||
|
||||
|
||||
class SpecialistExecutor(CrewAIBaseSpecialistExecutor):
|
||||
"""
|
||||
type: SPIN_SPECIALIST
|
||||
type_version: 1.0
|
||||
SPIN Specialist Executor class
|
||||
"""
|
||||
|
||||
def __init__(self, tenant_id, specialist_id, session_id, task_id, **kwargs):
|
||||
self.rag_crew = None
|
||||
self.spin_crew = None
|
||||
self.identification_crew = None
|
||||
self.rag_consolidation_crew = None
|
||||
|
||||
super().__init__(tenant_id, specialist_id, session_id, task_id)
|
||||
|
||||
# Load the Tenant & set language
|
||||
self.tenant = Tenant.query.get_or_404(tenant_id)
|
||||
if self.specialist.configuration['tenant_language'] is None:
|
||||
self.specialist.configuration['tenant_language'] = self.tenant.language
|
||||
|
||||
@property
|
||||
def type(self) -> str:
|
||||
return "SPIN_SPECIALIST"
|
||||
|
||||
@property
|
||||
def type_version(self) -> str:
|
||||
return "1.0"
|
||||
|
||||
def _config_task_agents(self):
|
||||
self._add_task_agent("rag_task", "rag_agent")
|
||||
self._add_task_agent("spin_detect_task", "spin_detection_agent")
|
||||
self._add_task_agent("spin_questions_task", "spin_sales_specialist_agent")
|
||||
self._add_task_agent("identification_detection_task", "identification_agent")
|
||||
self._add_task_agent("identification_questions_task", "identification_agent")
|
||||
self._add_task_agent("email_lead_drafting_task", "email_content_agent")
|
||||
self._add_task_agent("email_lead_engagement_task", "email_engagement_agent")
|
||||
self._add_task_agent("email_lead_retrieval_task", "email_engagement_agent")
|
||||
self._add_task_agent("rag_consolidation_task", "rag_communication_agent")
|
||||
|
||||
def _config_pydantic_outputs(self):
|
||||
self._add_pydantic_output("rag_task", RAGOutput, "rag_output")
|
||||
self._add_pydantic_output("spin_questions_task", SPINOutput, "spin_questions")
|
||||
self._add_pydantic_output("identification_questions_task", LeadInfoOutput, "lead_identification_questions")
|
||||
self._add_pydantic_output("rag_consolidation_task", RAGOutput, "rag_output")
|
||||
|
||||
def _instantiate_specialist(self):
|
||||
verbose = self.tuning
|
||||
|
||||
rag_agents = [self.rag_agent]
|
||||
rag_tasks = [self.rag_task]
|
||||
self.rag_crew = EveAICrewAICrew(
|
||||
self,
|
||||
"Rag Crew",
|
||||
agents=rag_agents,
|
||||
tasks=rag_tasks,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
spin_agents = [self.spin_detection_agent, self.spin_sales_specialist_agent]
|
||||
spin_tasks = [self.spin_detect_task, self.spin_questions_task]
|
||||
self.spin_crew = EveAICrewAICrew(
|
||||
self,
|
||||
"SPIN Crew",
|
||||
agents=spin_agents,
|
||||
tasks=spin_tasks,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
identification_agents = [self.identification_agent]
|
||||
identification_tasks = [self.identification_detection_task, self.identification_questions_task]
|
||||
self.identification_crew = EveAICrewAICrew(
|
||||
self,
|
||||
"Identification Crew",
|
||||
agents=identification_agents,
|
||||
tasks=identification_tasks,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
consolidation_agents = [self.rag_communication_agent]
|
||||
consolidation_tasks = [self.rag_consolidation_task]
|
||||
self.rag_consolidation_crew = EveAICrewAICrew(
|
||||
self,
|
||||
"Rag Consolidation Crew",
|
||||
agents=consolidation_agents,
|
||||
tasks=consolidation_tasks,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
self.flow = SPINFlow(
|
||||
self,
|
||||
self.rag_crew,
|
||||
self.spin_crew,
|
||||
self.identification_crew,
|
||||
self.rag_consolidation_crew
|
||||
)
|
||||
|
||||
def execute(self, arguments: SpecialistArguments) -> SpecialistResult:
|
||||
formatted_context, citations = self.retrieve_context(arguments)
|
||||
|
||||
self.log_tuning("SPIN Specialist execution started", {})
|
||||
|
||||
flow_inputs = {
|
||||
"language": arguments.language,
|
||||
"query": arguments.query,
|
||||
"context": formatted_context,
|
||||
"citations": citations,
|
||||
"history": self._formatted_history,
|
||||
"name": self.specialist.configuration.get('name', ''),
|
||||
"company": self.specialist.configuration.get('company', ''),
|
||||
"products": self.specialist.configuration.get('products', ''),
|
||||
"product_information": self.specialist.configuration.get('product_information', ''),
|
||||
"engagement_options": self.specialist.configuration.get('engagement_options', ''),
|
||||
"tenant_language": self.specialist.configuration.get('tenant_language', ''),
|
||||
"nr_of_questions": self.specialist.configuration.get('nr_of_questions', ''),
|
||||
}
|
||||
# crew_results = self.rag_crew.kickoff(inputs=flow_inputs)
|
||||
# current_app.logger.debug(f"Test Crew Output received: {crew_results}")
|
||||
flow_results = self.flow.kickoff(inputs=flow_inputs)
|
||||
|
||||
flow_state = self.flow.state
|
||||
|
||||
results = SPINSpecialistResult.create_for_type(self.type, self.type_version)
|
||||
update_data = {}
|
||||
if flow_state.final_output:
|
||||
update_data["rag_output"] = flow_state.final_output
|
||||
elif flow_state.rag_output: # Fallback
|
||||
update_data["rag_output"] = flow_state.rag_output
|
||||
if flow_state.spin:
|
||||
update_data["spin"] = flow_state.spin
|
||||
if flow_state.lead_info:
|
||||
update_data["lead_info"] = flow_state.lead_info
|
||||
|
||||
results = results.model_copy(update=update_data)
|
||||
|
||||
self.log_tuning(f"SPIN Specialist execution ended", {"Results": results.model_dump()})
|
||||
|
||||
return results
|
||||
|
||||
# TODO: metrics
|
||||
|
||||
|
||||
class SPINSpecialistInput(BaseModel):
|
||||
language: Optional[str] = Field(None, alias="language")
|
||||
query: Optional[str] = Field(None, alias="query")
|
||||
context: Optional[str] = Field(None, alias="context")
|
||||
citations: Optional[List[int]] = Field(None, alias="citations")
|
||||
history: Optional[str] = Field(None, alias="history")
|
||||
name: Optional[str] = Field(None, alias="name")
|
||||
company: Optional[str] = Field(None, alias="company")
|
||||
products: Optional[str] = Field(None, alias="products")
|
||||
product_information: Optional[str] = Field(None, alias="product_information")
|
||||
engagement_options: Optional[str] = Field(None, alias="engagement_options")
|
||||
tenant_language: Optional[str] = Field(None, alias="tenant_language")
|
||||
nr_of_questions: Optional[int] = Field(None, alias="nr_of_questions")
|
||||
|
||||
|
||||
class SPINSpecialistResult(SpecialistResult):
|
||||
rag_output: Optional[RAGOutput] = Field(None, alias="Rag Output")
|
||||
spin: Optional[SPINOutput] = Field(None, alias="Spin Output")
|
||||
lead_info: Optional[LeadInfoOutput] = Field(None, alias="Lead Info Output")
|
||||
|
||||
|
||||
class SPINFlowState(EveAIFlowState):
|
||||
"""Flow state for SPIN specialist that automatically updates from task outputs"""
|
||||
input: Optional[SPINSpecialistInput] = None
|
||||
rag_output: Optional[RAGOutput] = None
|
||||
lead_info: Optional[LeadInfoOutput] = None
|
||||
spin: Optional[SPINOutput] = None
|
||||
final_output: Optional[RAGOutput] = None
|
||||
|
||||
|
||||
class SPINFlow(EveAICrewAIFlow[SPINFlowState]):
|
||||
def __init__(self, specialist_executor, rag_crew, spin_crew, identification_crew, rag_consolidation_crew, **kwargs):
|
||||
super().__init__(specialist_executor, "SPIN Specialist Flow", **kwargs)
|
||||
self.specialist_executor = specialist_executor
|
||||
self.rag_crew = rag_crew
|
||||
self.spin_crew = spin_crew
|
||||
self.identification_crew = identification_crew
|
||||
self.rag_consolidation_crew = rag_consolidation_crew
|
||||
self.exception_raised = False
|
||||
|
||||
@start()
|
||||
def process_inputs(self):
|
||||
return ""
|
||||
|
||||
@listen(process_inputs)
|
||||
def execute_rag(self):
|
||||
inputs = self.state.input.model_dump()
|
||||
try:
|
||||
crew_output = self.rag_crew.kickoff(inputs=inputs)
|
||||
self.specialist_executor.log_tuning("RAG Crew Output", crew_output.model_dump())
|
||||
output_pydantic = crew_output.pydantic
|
||||
if not output_pydantic:
|
||||
raw_json = json.loads(crew_output.raw)
|
||||
output_pydantic = RAGOutput.model_validate(raw_json)
|
||||
self.state.rag_output = output_pydantic
|
||||
return crew_output
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"CREW rag_crew Kickoff Error: {str(e)}")
|
||||
self.exception_raised = True
|
||||
raise e
|
||||
|
||||
@listen(process_inputs)
|
||||
def execute_spin(self):
|
||||
inputs = self.state.input.model_dump()
|
||||
try:
|
||||
crew_output = self.spin_crew.kickoff(inputs=inputs)
|
||||
self.specialist_executor.log_tuning("Spin Crew Output", crew_output.model_dump())
|
||||
output_pydantic = crew_output.pydantic
|
||||
if not output_pydantic:
|
||||
raw_json = json.loads(crew_output.raw)
|
||||
output_pydantic = SPINOutput.model_validate(raw_json)
|
||||
self.state.spin = output_pydantic
|
||||
return crew_output
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"CREW spin_crew Kickoff Error: {str(e)}")
|
||||
self.exception_raised = True
|
||||
raise e
|
||||
|
||||
@listen(process_inputs)
|
||||
def execute_identification(self):
|
||||
inputs = self.state.input.model_dump()
|
||||
try:
|
||||
crew_output = self.identification_crew.kickoff(inputs=inputs)
|
||||
self.specialist_executor.log_tuning("Identification Crew Output", crew_output.model_dump())
|
||||
output_pydantic = crew_output.pydantic
|
||||
if not output_pydantic:
|
||||
raw_json = json.loads(crew_output.raw)
|
||||
output_pydantic = LeadInfoOutput.model_validate(raw_json)
|
||||
self.state.lead_info = output_pydantic
|
||||
return crew_output
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"CREW identification_crew Kickoff Error: {str(e)}")
|
||||
self.exception_raised = True
|
||||
raise e
|
||||
|
||||
@listen(and_(execute_rag, execute_spin, execute_identification))
|
||||
def consolidate(self):
|
||||
inputs = self.state.input.model_dump()
|
||||
if self.state.rag_output:
|
||||
inputs["prepared_answers"] = self.state.rag_output.answer
|
||||
additional_questions = ""
|
||||
if self.state.lead_info:
|
||||
additional_questions = self.state.lead_info.questions + "\n"
|
||||
if self.state.spin:
|
||||
additional_questions = additional_questions + self.state.spin.questions
|
||||
inputs["additional_questions"] = additional_questions
|
||||
try:
|
||||
crew_output = self.rag_consolidation_crew.kickoff(inputs=inputs)
|
||||
self.specialist_executor.log_tuning("RAG Consolidation Crew Output", crew_output.model_dump())
|
||||
output_pydantic = crew_output.pydantic
|
||||
if not output_pydantic:
|
||||
raw_json = json.loads(crew_output.raw)
|
||||
output_pydantic = LeadInfoOutput.model_validate(raw_json)
|
||||
self.state.final_output = output_pydantic
|
||||
return crew_output
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"CREW rag_consolidation_crew Kickoff Error: {str(e)}")
|
||||
self.exception_raised = True
|
||||
raise e
|
||||
|
||||
def kickoff(self, inputs=None):
|
||||
with current_event.create_span("SPIN Specialist Execution"):
|
||||
self.specialist_executor.log_tuning("Inputs retrieved", inputs)
|
||||
self.state.input = SPINSpecialistInput.model_validate(inputs)
|
||||
self.specialist.update_progress("EveAI Flow Start", {"name": "SPIN"})
|
||||
try:
|
||||
result = super().kickoff()
|
||||
except Exception as e:
|
||||
current_app.logger.error(f"Error kicking of Flow: {str(e)}")
|
||||
|
||||
self.specialist.update_progress("EveAI Flow End", {"name": "SPIN"})
|
||||
|
||||
return self.state
|
||||
Reference in New Issue
Block a user