dapr-agents/cookbook/workflows/llm/multi_modal_task_chain.py

69 lines
1.8 KiB
Python

from dapr_agents import OpenAIChatClient, NVIDIAChatClient
from dapr.ext.workflow import DaprWorkflowContext
from dapr_agents.workflow import WorkflowApp, task, workflow
from dotenv import load_dotenv
import os
import logging
load_dotenv()
nvidia_llm = NVIDIAChatClient(
model="meta/llama-3.1-8b-instruct", api_key=os.getenv("NVIDIA_API_KEY")
)
oai_llm = OpenAIChatClient(
api_key=os.getenv("OPENAI_API_KEY"),
model="gpt-4o",
base_url=os.getenv("OPENAI_API_BASE_URL"),
)
azoai_llm = OpenAIChatClient(
api_key=os.getenv("AZURE_OPENAI_API_KEY"),
azure_deployment="gpt-4o-mini",
azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
azure_api_version="2024-12-01-preview",
)
@workflow
def test_workflow(ctx: DaprWorkflowContext):
"""
A simple workflow that uses a multi-modal task chain.
"""
oai_results = yield ctx.call_activity(invoke_oai, input="Peru")
azoai_results = yield ctx.call_activity(invoke_azoai, input=oai_results)
nvidia_results = yield ctx.call_activity(invoke_nvidia, input=azoai_results)
return nvidia_results
@task(
description="What is the name of the capital of {country}?. Reply with just the name.",
llm=oai_llm,
)
def invoke_oai(country: str) -> str:
pass
@task(description="What is a famous thing about {capital}?", llm=azoai_llm)
def invoke_azoai(capital: str) -> str:
pass
@task(
description="Context: {context}. From the previous context. Pick one thing to do.",
llm=nvidia_llm,
)
def invoke_nvidia(context: str) -> str:
pass
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
wfapp = WorkflowApp()
results = wfapp.run_and_monitor_workflow_sync(workflow=test_workflow)
logging.info("Workflow results: %s", results)
logging.info("Workflow completed successfully.")