From 3e44263e0d6ed4512298f9142d273f80f326ccbf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Sun, 24 May 2026 20:10:21 +0000 Subject: [PATCH] Add agent_stream.py --- agent_stream.py | 90 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 agent_stream.py diff --git a/agent_stream.py b/agent_stream.py new file mode 100644 index 0000000..b5316f6 --- /dev/null +++ b/agent_stream.py @@ -0,0 +1,90 @@ +"""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()