Files
temporal-ai-agent/activities/tool_activities.py
2025-01-01 13:54:54 -08:00

128 lines
3.8 KiB
Python

from dataclasses import dataclass
from temporalio import activity
from temporalio.exceptions import ApplicationError
from ollama import chat, ChatResponse
import json
from models.tool_definitions import ToolsData
from typing import Sequence
from temporalio.common import RawValue
@dataclass
class ToolPromptInput:
prompt: str
context_instructions: str
class ToolActivities:
@activity.defn
def prompt_llm(self, input: ToolPromptInput) -> str:
model_name = "qwen2.5:14b"
messages = [
{
"role": "system",
"content": input.context_instructions
+ ". The current date is "
+ get_current_date_human_readable(),
},
{
"role": "user",
"content": input.prompt,
},
]
response: ChatResponse = chat(model=model_name, messages=messages)
return response.message.content
@activity.defn
def parse_tool_data(self, json_str: str) -> dict:
"""
Parses a JSON string into a dictionary.
Raises a ValueError if the JSON is invalid.
"""
try:
data = json.loads(json_str)
except json.JSONDecodeError as e:
raise ApplicationError(f"Invalid JSON: {e}")
return data
@activity.defn
def validate_and_parse_json(
self,
response_prechecked: str,
tools_data: ToolsData,
conversation_history: str,
) -> dict:
"""
1) Build JSON validation instructions
2) Call LLM with those instructions
3) Parse the result
4) If parsing fails, raise exception -> triggers retry
"""
# 1) Build validation instructions
# (Generate the validation prompt exactly as you do in your workflow.)
from prompts.agent_prompt_generators import (
generate_json_validation_prompt_from_tools_data,
)
validation_prompt = generate_json_validation_prompt_from_tools_data(
tools_data, conversation_history, response_prechecked
)
# 2) Call LLM
prompt_input = ToolPromptInput(
prompt=response_prechecked,
context_instructions=validation_prompt,
)
validated_response = self.prompt_llm(prompt_input)
# 3) Parse
# If parse fails, we raise ApplicationError -> triggers retry
try:
parsed = self.parse_tool_data(validated_response)
except Exception as e:
raise ApplicationError(f"Failed to parse validated JSON: {e}")
# 4) If we get here, parse succeeded
return parsed
def get_current_date_human_readable():
"""
Returns the current date in a human-readable format.
Example: Wednesday, January 1, 2025
"""
from datetime import datetime
return datetime.now().strftime("%A, %B %d, %Y")
@activity.defn(dynamic=True)
def dynamic_tool_activity(args: Sequence[RawValue]) -> dict:
"""Dynamic activity that is invoked via an unknown activity type."""
tool_name = activity.info().activity_type # e.g. "SearchFlights"
# The first payload is the dictionary of arguments
tool_args = activity.payload_converter().from_payload(args[0].payload, dict)
# Extract fields from the arguments
date_depart = tool_args.get("dateDepart")
date_return = tool_args.get("dateReturn")
origin = tool_args.get("origin")
destination = tool_args.get("destination")
# Print (or log) them
activity.logger.info(f"Tool: {tool_name}")
activity.logger.info(f"Depart: {date_depart}, Return: {date_return}")
activity.logger.info(f"Origin: {origin}, Destination: {destination}")
# For now, just return them
return {
"tool": tool_name,
"args": tool_args,
"status": "OK - dynamic activity stub",
}