Upload fluxai.py

fluxai.py

@@ -9,12 +9,16 @@ from dotenv import load_dotenv
 from pydantic import BaseModel
 from pymongo import MongoClient
 from models import *
+from huggingface_hub import InferenceClient
 
 class FluxAI(BaseModel):
     user_id: int
     args: str
     auto_enhancer: bool = False
 
+class MistralAI(BaseModel):
+    args: str
+
 router = APIRouter()
 
 load_dotenv()
@@ -35,6 +39,20 @@ async def schellwithflux(args):
         return None
     return response.content
 
+async def mistralai_post_message(message_str):
+    client = InferenceClient(
+        "mistralai/Mixtral-8x7B-Instruct-v0.1",
+        token=HUGGING_TOKEN
+    )
+    output = ""
+    for message in client.chat_completion(
+        messages=[{"role": "user", "content": message_str}],
+        max_tokens=500,
+        stream=True
+    ):
+        output += message.choices[0].delta.content
+    return output
+
 def get_user_tokens_gpt(user_id):
     user = collection.find_one({"user_id": user_id})
     if not user:
@@ -52,6 +70,20 @@ def deduct_tokens_gpt(user_id, amount):
     else:
         return False
 
+@router.post("/akeno/mistralai", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
+async def mistralai_(payload: MistralAI):
+        try:
+            response = await mistralai_post_message(payload.args)
+            return SuccessResponse(
+                status="True",
+                randydev={"message": response}
+            )
+        except Exception as e:
+            return SuccessResponse(
+                status="False",
+                randydev={"error": f"An error occurred: {str(e)}"}
+            )
+
 @router.post("/akeno/fluxai", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
 async def fluxai_image(payload: FluxAI):
     if deduct_tokens_gpt(payload.user_id, amount=20):