Add stream_agent.py
This commit is contained in:
@@ -0,0 +1,77 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
class StreamingAIAgent:
|
||||
"""Small AI-agent with a streaming interface and an offline fallback."""
|
||||
|
||||
def __init__(self, model: str | None = None) -> None:
|
||||
self.model = model or os.getenv("OPENAI_MODEL", "openai/gpt-oss-20b:free")
|
||||
|
||||
def invoke(self, question: str) -> str:
|
||||
return "".join(self.stream(question))
|
||||
|
||||
def stream(self, question: str) -> Iterable[str]:
|
||||
client = self._build_client()
|
||||
if client is None:
|
||||
yield from self._offline_stream(question)
|
||||
return
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a concise educational assistant.",
|
||||
},
|
||||
{"role": "user", "content": question},
|
||||
],
|
||||
stream=True,
|
||||
)
|
||||
for chunk in response:
|
||||
delta = chunk.choices[0].delta.content
|
||||
if delta:
|
||||
yield delta
|
||||
|
||||
def _build_client(self):
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
if not api_key:
|
||||
return None
|
||||
try:
|
||||
from openai import OpenAI
|
||||
except ImportError:
|
||||
return None
|
||||
return OpenAI(
|
||||
api_key=api_key,
|
||||
base_url=os.getenv("OPENAI_BASE_URL") or None,
|
||||
)
|
||||
|
||||
def _offline_stream(self, question: str) -> Iterable[str]:
|
||||
answer = (
|
||||
"Stream mode returns an answer piece by piece instead of waiting "
|
||||
"for the whole response. This makes an AI-agent feel faster and "
|
||||
f"lets the UI show progress while processing: {question}"
|
||||
)
|
||||
for token in answer.split(" "):
|
||||
yield token + " "
|
||||
time.sleep(0.02)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Streaming AI-agent demo")
|
||||
parser.add_argument("question", nargs="*", help="User question")
|
||||
args = parser.parse_args()
|
||||
|
||||
question = " ".join(args.question).strip() or "Explain streaming mode"
|
||||
agent = StreamingAIAgent()
|
||||
for token in agent.stream(question):
|
||||
print(token, end="", flush=True)
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user