68 lines
2.5 KiB
Python
68 lines
2.5 KiB
Python
from __future__ import annotations
|
|
from typing import Literal, AsyncGenerator, ClassVar, Optional
|
|
import logging
|
|
from . base import Agent, registry
|
|
from .. message import Message
|
|
import inspect
|
|
|
|
class Chat(Agent):
|
|
"""
|
|
Chat Agent
|
|
"""
|
|
agent_type: Literal["chat"] = "chat" # type: ignore
|
|
_agent_type: ClassVar[str] = agent_type # Add this for registration
|
|
|
|
async def prepare_message(self, message:Message) -> AsyncGenerator[Message, None]:
|
|
"""
|
|
Prepare message with context information in message.preamble
|
|
"""
|
|
logging.info(f"{self.agent_type} - {inspect.stack()[1].function}")
|
|
|
|
if not self.context:
|
|
raise ValueError("Context is not set for this agent.")
|
|
|
|
# Generate RAG content if enabled, based on the content
|
|
rag_context = ""
|
|
if message.enable_rag:
|
|
# Gather RAG results, yielding each result
|
|
# as it becomes available
|
|
for message in self.context.generate_rag_results(message):
|
|
logging.info(f"RAG: {message.status} - {message.response}")
|
|
if message.status == "error":
|
|
yield message
|
|
return
|
|
if message.status != "done":
|
|
yield message
|
|
|
|
if "rag" in message.metadata and message.metadata["rag"]:
|
|
for rag in message.metadata["rag"]:
|
|
for doc in rag["documents"]:
|
|
rag_context += f"{doc}\n"
|
|
|
|
message.preamble = {}
|
|
|
|
if rag_context:
|
|
message.preamble["context"] = rag_context
|
|
|
|
if self.context.user_resume:
|
|
message.preamble["resume"] = self.context.user_resume
|
|
|
|
if message.preamble:
|
|
preamble_types = [f"<|{p}|>" for p in message.preamble.keys()]
|
|
preamble_types_AND = " and ".join(preamble_types)
|
|
preamble_types_OR = " or ".join(preamble_types)
|
|
message.preamble["rules"] = f"""\
|
|
- Answer the question based on the information provided in the {preamble_types_AND} sections by incorporate it seamlessly and refer to it using natural language instead of mentioning {preamble_types_OR} or quoting it directly.
|
|
- If there is no information in these sections, answer based on your knowledge, or use any available tools.
|
|
- Avoid phrases like 'According to the {preamble_types[0]}' or similar references to the {preamble_types_OR}.
|
|
"""
|
|
message.preamble["question"] = "Respond to:"
|
|
|
|
message.system_prompt = self.system_prompt
|
|
message.status = "done"
|
|
yield message
|
|
return
|
|
|
|
# Register the base agent
|
|
registry.register(Chat._agent_type, Chat)
|