Update app.py with mistralai streaming api

app.py

@@ -1,7 +1,36 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Depends
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Annotated
+from mistralai import Mistral
+import os
 
 app = FastAPI()
 
+mistral = os.environ.get('MISTRAL_KEY', '')
+mistral_client = Mistral(api_key=mistral)
+
 @app.get("/")
 def hello():
-    return {"Hello": "World!"}
+    return {"Hello": "World!"}
+
+class MistralRequest(BaseModel):
+    model: str
+    prompt: str
+
+@app.post("/mistral")
+async def mistral(request: MistralRequest, token: Annotated[str, Depends(verify_token)] 
+    async def generate():
+        response = await mistral_client.chat.stream_async(
+            model=request.model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": request.prompt,
+                }
+            ],
+        )
+        async for chunk in response:
+            if chunk.data.choices[0].delta.content is not None:
+                yield chunk.data.choices[0].delta.content
+
+    return StreamingResponse(generate(), media_type="text/plain")