78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
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()
|