91 lines
2.5 KiB
Python
91 lines
2.5 KiB
Python
"""Simple streaming AI agent using OpenAI ChatCompletion API.
|
||
|
||
The module exposes a single function ``stream_chat`` that yields
|
||
text fragments as they arrive from the model. It is intentionally
|
||
light‑weight so it can be used as a drop‑in component in larger
|
||
applications.
|
||
|
||
Dependencies
|
||
------------
|
||
* ``openai`` – official OpenAI SDK
|
||
|
||
Environment
|
||
-----------
|
||
An OpenAI API key must be available either as the environment
|
||
variable ``OPENAI_API_KEY`` or passed explicitly via the
|
||
``api_key`` argument.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from typing import Iterable, Generator, Dict, List, Any
|
||
|
||
import openai
|
||
|
||
# Ensure the OpenAI key is set when the module is imported.
|
||
# Users can override by passing ``api_key`` to ``stream_chat``.
|
||
openai.api_key = os.getenv("OPENAI_API_KEY")
|
||
|
||
|
||
def stream_chat(
|
||
messages: List[Dict[str, str]],
|
||
*,
|
||
model: str = "gpt-3.5-turbo",
|
||
temperature: float = 0.7,
|
||
api_key: str | None = None,
|
||
) -> Generator[str, None, None]:
|
||
"""Yield model output token by token.
|
||
|
||
Parameters
|
||
----------
|
||
messages:
|
||
A list of message objects compatible with the ChatCompletion API.
|
||
model:
|
||
The model to use. Defaults to ``gpt-3.5-turbo``.
|
||
temperature:
|
||
Sampling temperature. Defaults to 0.7.
|
||
api_key:
|
||
Optional API key. If provided it overrides the environment
|
||
variable.
|
||
|
||
Yields
|
||
------
|
||
str
|
||
The next fragment of the assistant's reply.
|
||
"""
|
||
|
||
if api_key is not None:
|
||
openai.api_key = api_key
|
||
|
||
# The SDK returns an iterator over chunks when stream=True.
|
||
response = openai.ChatCompletion.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
stream=True,
|
||
)
|
||
|
||
# The response is an iterator of chunks. Each chunk contains a
|
||
# ``choices[0].delta`` dict with the text fragment.
|
||
for chunk in response:
|
||
try:
|
||
delta = chunk["choices"][0]["delta"]
|
||
if "content" in delta:
|
||
yield delta["content"]
|
||
except Exception as exc: # pragma: no cover – defensive
|
||
# In a real application you might log this.
|
||
print(f"Error processing chunk: {exc}")
|
||
continue
|
||
|
||
|
||
# If the module is executed directly, run a small demo.
|
||
if __name__ == "__main__": # pragma: no cover
|
||
demo_messages = [
|
||
{"role": "user", "content": "Write a short poem about the ocean."}
|
||
]
|
||
print("Streaming response: ")
|
||
for token in stream_chat(demo_messages):
|
||
print(token, end="", flush=True)
|
||
print()
|