72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
from pydantic import BaseModel, Field
|
|
from typing import Dict, List, Optional, Any, Union, Mapping
|
|
from datetime import datetime, timezone
|
|
from . rag import ChromaDBGetResponse
|
|
from ollama._types import Options
|
|
|
|
class Tunables(BaseModel):
|
|
enable_rag: bool = True # Enable RAG collection chromadb matching
|
|
enable_tools: bool = True # Enable LLM to use tools
|
|
enable_context: bool = True # Add <|context|> field to message
|
|
|
|
class MessageMetaData(BaseModel):
|
|
rag: List[ChromaDBGetResponse] = Field(default_factory=list)
|
|
eval_count: int = 0
|
|
eval_duration: int = 0
|
|
prompt_eval_count: int = 0
|
|
prompt_eval_duration: int = 0
|
|
context_size: int = 0
|
|
resume_generation: Optional[Dict[str, Any]] = None
|
|
options: Optional[Union[Mapping[str, Any], Options]] = None
|
|
tools: Optional[Dict[str, Any]] = None
|
|
timers: Optional[Dict[str, float]] = None
|
|
|
|
#resume : str = ""
|
|
#match_stats: Optional[Dict[str, Dict[str, Any]]] = Field(default=None)
|
|
|
|
class Message(BaseModel):
|
|
model_config = {"arbitrary_types_allowed": True} # Allow Event
|
|
# Required
|
|
prompt: str # Query to be answered
|
|
|
|
# Tunables
|
|
tunables: Tunables = Field(default_factory=Tunables)
|
|
|
|
# Generated while processing message
|
|
status: str = "" # Status of the message
|
|
preamble: Dict[str, Any] = Field(default_factory=dict) # Preamble to be prepended to the prompt
|
|
system_prompt: str = "" # System prompt provided to the LLM
|
|
context_prompt: str = "" # Full content of the message (preamble + prompt)
|
|
response: str = "" # LLM response to the preamble + query
|
|
metadata: MessageMetaData = Field(default_factory=MessageMetaData)
|
|
network_packets: int = 0 # Total number of streaming packets
|
|
network_bytes: int = 0 # Total bytes sent while streaming packets
|
|
actions: List[str] = (
|
|
[]
|
|
) # Other session modifying actions performed while processing the message
|
|
timestamp: str = str(datetime.now(timezone.utc))
|
|
chunk: str = ""
|
|
partial_response: str = ""
|
|
title: str = ""
|
|
context_size: int = 0
|
|
|
|
def add_action(self, action: str | list[str]) -> None:
|
|
"""Add a actions(s) to the message."""
|
|
if isinstance(action, str):
|
|
self.actions.append(action)
|
|
else:
|
|
self.actions.extend(action)
|
|
|
|
def get_summary(self) -> str:
|
|
"""Return a summary of the message."""
|
|
response_summary = (
|
|
f"Response: {self.response} (Actions: {', '.join(self.actions)})"
|
|
if self.response
|
|
else "No response yet"
|
|
)
|
|
return (
|
|
f"Message at {self.timestamp}:\n"
|
|
f"Query: {self.preamble}{self.content}\n"
|
|
f"{response_summary}"
|
|
)
|