Add agent_stream.py

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()