fork download
  1. import asyncio
  2. import os
  3. from typing import AsyncIterator
  4. from uuid import uuid4
  5.  
  6. from literalai import AsyncLiteralClient
  7. from openai import AsyncOpenAI
  8.  
  9. LITERAL_API_KEY = os.getenv("LITERAL_API_KEY")
  10. OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
  11. PROMPT = """Return the user message uppercase"""
  12.  
  13. literalai_client = AsyncLiteralClient(api_key=LITERAL_API_KEY, environment="dev")
  14. literalai_client.initialize()
  15. openai_client = AsyncOpenAI(api_key=OPENAI_API_KEY)
  16.  
  17.  
  18. async def call_completion(msg: str) -> AsyncIterator[str]:
  19. stream = await openai_client.chat.completions.create(
  20. model="gpt-4o-mini",
  21. messages=[
  22. {"role": "system", "content": PROMPT},
  23. {"role": "user", "content": msg},
  24. ],
  25. stream=True,
  26. )
  27.  
  28. async for chunk in stream:
  29. if chunk.choices[0].delta.content is not None:
  30. yield chunk.choices[0].delta.content
  31.  
  32.  
  33. async def make_request(num: int) -> None:
  34. print(f"Running {num}")
  35. thread_id = f"req{num}_{uuid4()}"
  36. msg = f"Hello {num}"
  37.  
  38. with literalai_client.thread(thread_id=thread_id, name=thread_id):
  39. literalai_client.message(content=msg, type="user_message")
  40. print(f"Querying {num}")
  41. response = ""
  42. with literalai_client.step(type="run", name="run_stream"):
  43. async for chunk in call_completion(msg):
  44. response += chunk
  45. await asyncio.sleep(0)
  46.  
  47. print(f"Got response {num}")
  48. literalai_client.message(content=response, type="assistant_message")
  49.  
  50.  
  51. async def main() -> None:
  52. requests = [make_request(i) for i in range(3)]
  53. await asyncio.gather(*requests)
  54. literalai_client.flush_and_stop()
  55.  
  56.  
  57. if __name__ == "__main__":
  58. asyncio.run(main())
  59.  
Runtime error #stdin #stdout #stderr 1.3s 38688KB
stdin
Standard input is empty
stdout
Standard output is empty
stderr
Traceback (most recent call last):
  File "./prog.py", line 6, in <module>
ModuleNotFoundError: No module named 'literalai'