|
| 1 | +import json |
| 2 | +import os |
| 3 | + |
| 4 | +import openai |
| 5 | +from openai.types.responses import ResponseInputParam |
| 6 | +from slack_sdk.models.messages.chunk import TaskUpdateChunk |
| 7 | +from slack_sdk.web.chat_stream import ChatStream |
| 8 | + |
| 9 | +from agent.tools.dice import roll_dice, roll_dice_definition |
| 10 | + |
| 11 | + |
| 12 | +def call_llm( |
| 13 | + streamer: ChatStream, |
| 14 | + prompts: ResponseInputParam, |
| 15 | +): |
| 16 | + """ |
| 17 | + Stream an LLM response to prompts with an example dice rolling function |
| 18 | +
|
| 19 | + https://docs.slack.dev/tools/python-slack-sdk/web#sending-streaming-messages |
| 20 | + https://platform.openai.com/docs/guides/text |
| 21 | + https://platform.openai.com/docs/guides/streaming-responses |
| 22 | + https://platform.openai.com/docs/guides/function-calling |
| 23 | + """ |
| 24 | + llm = openai.OpenAI( |
| 25 | + api_key=os.getenv("OPENAI_API_KEY"), |
| 26 | + ) |
| 27 | + tool_calls = [] |
| 28 | + response = llm.responses.create( |
| 29 | + model="gpt-4o-mini", |
| 30 | + input=prompts, |
| 31 | + tools=[ |
| 32 | + roll_dice_definition, |
| 33 | + ], |
| 34 | + stream=True, |
| 35 | + ) |
| 36 | + for event in response: |
| 37 | + # Markdown text from the LLM response is streamed in chat as it arrives |
| 38 | + if event.type == "response.output_text.delta": |
| 39 | + streamer.append(markdown_text=f"{event.delta}") |
| 40 | + |
| 41 | + # Function calls are saved for later computation and a new task is shown |
| 42 | + if event.type == "response.output_item.done": |
| 43 | + if event.item.type == "function_call": |
| 44 | + tool_calls.append(event.item) |
| 45 | + if event.item.name == "roll_dice": |
| 46 | + args = json.loads(event.item.arguments) |
| 47 | + streamer.append( |
| 48 | + chunks=[ |
| 49 | + TaskUpdateChunk( |
| 50 | + id=f"{event.item.call_id}", |
| 51 | + title=f"Rolling a {args['count']}d{args['sides']}...", |
| 52 | + status="in_progress", |
| 53 | + ), |
| 54 | + ], |
| 55 | + ) |
| 56 | + |
| 57 | + # Tool calls are performed and tasks are marked as completed in Slack |
| 58 | + if tool_calls: |
| 59 | + for call in tool_calls: |
| 60 | + if call.name == "roll_dice": |
| 61 | + args = json.loads(call.arguments) |
| 62 | + prompts.append( |
| 63 | + { |
| 64 | + "id": call.id, |
| 65 | + "call_id": call.call_id, |
| 66 | + "type": "function_call", |
| 67 | + "name": "roll_dice", |
| 68 | + "arguments": call.arguments, |
| 69 | + } |
| 70 | + ) |
| 71 | + result = roll_dice(**args) |
| 72 | + prompts.append( |
| 73 | + { |
| 74 | + "type": "function_call_output", |
| 75 | + "call_id": call.call_id, |
| 76 | + "output": json.dumps(result), |
| 77 | + } |
| 78 | + ) |
| 79 | + if result.get("error") is not None: |
| 80 | + streamer.append( |
| 81 | + chunks=[ |
| 82 | + TaskUpdateChunk( |
| 83 | + id=f"{call.call_id}", |
| 84 | + title=f"{result['error']}", |
| 85 | + status="error", |
| 86 | + ), |
| 87 | + ], |
| 88 | + ) |
| 89 | + else: |
| 90 | + streamer.append( |
| 91 | + chunks=[ |
| 92 | + TaskUpdateChunk( |
| 93 | + id=f"{call.call_id}", |
| 94 | + title=f"{result['description']}", |
| 95 | + status="complete", |
| 96 | + ), |
| 97 | + ], |
| 98 | + ) |
| 99 | + |
| 100 | + # Complete the LLM response after making tool calls |
| 101 | + call_llm(streamer, prompts) |
0 commit comments