task-agent.py

from typing import Optional, List
from uuid import uuid4
from uagents import Agent, Context, Model, Field, Protocol
from uagents_core.contrib.protocols.chat import (
    ChatMessage,
    ChatAcknowledgement,
    TextContent,
    chat_protocol_spec
)
import requests
from datetime import datetime
from llm_output import llm_response
import os
from dotenv import load_dotenv
from llm_output import llm_response
load_dotenv()

OPEN_ROUTER_API_KEY = os.getenv("OPEN_ROUTER_API_KEY")

class RequestModel(Model):
    query: str = Field(...)

class ResponseModel(Model):
    tasks: list[str] = Field(...)
    error: Optional[str] = None

agent = Agent(name="task_agent", seed="task_secret_seed2")

@agent.on_event("startup")
async def startup_function(ctx: Context):
    ctx.logger.info(f"""Hello, I am an agent that extracts
                    actionable tasks from a message""")

def create_tasks_with_llm(query: str) -> list[str]:
    """Call LLM to turn a natural-language query into structured tasks."""
    # TODO: make prompt 
    prompt = f"""
    Extract a concise list of actionable tasks from this text. 
    Respond ONLY list of tasks, no explanation, separated by a '|' at the
    end of each task. Only show the tasks with no new lines, numbers, etc.

    Query: "{query}"
    """

    text = llm_response(OPEN_ROUTER_API_KEY, prompt, 400)

    # Split the response into individual tasks
    tasks = [line.strip() for line in text.split("|") if line.strip()]
    return ResponseModel(
        tasks = tasks
    )


@agent.on_rest_post("/tasks", RequestModel, ResponseModel)
async def handle_query(ctx: Context, req: RequestModel) -> ResponseModel:
    try:
        query = req.query
        tasks_response = create_tasks_with_llm(query)
        ctx.logger.info(f"Query: {query} -> Tasks: {tasks_response.tasks}")
        return tasks_response
    except Exception as e:
        return ResponseModel(
            tasks=[],
            error="No tasks found"
        )

protocol = Protocol(spec=chat_protocol_spec)

@protocol.on_message(ChatMessage)
async def handle_message(ctx: Context, sender: str, msg: ChatMessage):
    await ctx.send(
        sender,
        ChatAcknowledgement(timestamp=datetime.now(), acknowledged_msg_id=msg.msg_id),
    )

    url = "http://127.0.0.1:8000/tasks"  # use 127.0.0.1 instead of 0.0.0.0 for local access

    # Example request body based on your RequestModel
    payload = {
        "query": "How do I fine-tune a transformer model?",  # adjust field names to match your RequestModel
    }


    response: ResponseModel = requests.post(url, json=payload)
    task_list: List[str] = response.tasks

    is_user_message = isinstance(msg.content[-1], TextContent)
    ctx.logger.info(msg)
    if is_user_message:
        await ctx.send(sender, ChatMessage(
            timestamp=datetime.now(),
            msg_id=uuid4(),
            content=[
                # TODO: no hard code
                TextContent(type="text", text= " | ".join(task_list)),
                # This will end the chat session
                # EndSessionContent(type="end-session") 
            ]
        ))
@protocol.on_message(ChatAcknowledgement)
async def handle_ack(ctx: Context, sender: str, msg: ChatAcknowledgement):
    pass

agent.include(protocol, publish_manifest=True)
if __name__ == "__main__":
    agent.run()