Update app.py

app.py

@@ -5,9 +5,11 @@ from typing import List, Dict, Any, Optional
 from pydantic import BaseModel
 import asyncio
 import httpx
-
-from config import cookies, headers
+import random
+from config import cookies, headers, groqapi
 from prompts import ChiplingPrompts
+from groq import Groq
+import json
 
 app = FastAPI()
 
@@ -26,6 +28,8 @@ class ChatRequest(BaseModel):
     messages: List[Dict[Any, Any]]
     model: Optional[str] = "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8"
 
+client = Groq(api_key=groqapi)
+
 async def generate(json_data: Dict[str, Any]):
     max_retries = 5
     for attempt in range(max_retries):
@@ -60,6 +64,84 @@ async def generate(json_data: Dict[str, Any]):
 
     yield "data: [Max retries reached]\n\n"
 
+def convert_to_groq_schema(messages: List[Dict[str, Any]]) -> List[Dict[str, str]]:
+    converted = []
+    for message in messages:
+        role = message.get("role", "user")
+        content = message.get("content")
+
+        if isinstance(content, list):
+            flattened = []
+            for item in content:
+                if isinstance(item, dict) and item.get("type") == "text":
+                    flattened.append(item.get("text", ""))
+            content = "\n".join(flattened)
+        elif not isinstance(content, str):
+            content = str(content)
+
+        converted.append({"role": role, "content": content})
+    return converted
+
+
+async def groqgenerate(json_data: Dict[str, Any]):
+    try:
+        messages = convert_to_groq_schema(json_data["messages"])
+        chunk_id = "groq-" + "".join(random.choices("0123456789abcdef", k=32))
+        created = int(asyncio.get_event_loop().time())
+
+        # Create streaming response
+        stream = client.chat.completions.create(
+            messages=messages,
+            model="meta-llama/llama-4-scout-17b-16e-instruct",
+            temperature=json_data.get("temperature", 0.7),
+            max_completion_tokens=json_data.get("max_tokens", 1024),
+            top_p=json_data.get("top_p", 1),
+            stop=json_data.get("stop", None),
+            stream=True,
+        )
+
+        total_tokens = 0
+
+        # Use normal for-loop since stream is not async
+        for chunk in stream:
+            content = chunk.choices[0].delta.content
+            if content:
+                response = {
+                    "id": chunk_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": json_data.get("model", "llama-3.3-70b-versatile"),
+                    "choices": [{
+                        "index": 0,
+                        "text": content,
+                        "logprobs": None,
+                        "finish_reason": None
+                    }],
+                    "usage": None
+                }
+                yield f"data: {json.dumps(response)}\n\n"
+                total_tokens += 1
+
+        final = {
+            "id": chunk_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": json_data.get("model", "llama-3.3-70b-versatile"),
+            "choices": [],
+            "usage": {
+                "prompt_tokens": len(messages),
+                "completion_tokens": total_tokens,
+                "total_tokens": len(messages) + total_tokens,
+            }
+        }
+        yield f"data: {json.dumps(final)}\n\n"
+        yield "data: [DONE]\n\n"
+
+    except Exception as e:
+        return generate(json_data)
+
+
+
 @app.get("/")
 async def index():
     return {"status": "ok"}
@@ -93,7 +175,8 @@ async def chat(request: ChatRequest):
         'stream': True,
     }
 
-    return StreamingResponse(generate(json_data), media_type='text/event-stream')
+    selected_generator = random.choice([groqgenerate, generate])
+    return StreamingResponse(selected_generator(json_data), media_type='text/event-stream')
 
 
 @app.post("/generate-modules")
@@ -135,8 +218,8 @@ async def generate_modules(request: Request):
         'messages': current_messages,
         'stream': True,
     }
-
-    return StreamingResponse(generate(json_data), media_type='text/event-stream')
+    selected_generator = random.choice([groqgenerate])
+    return StreamingResponse(selected_generator(json_data), media_type='text/event-stream')
 
 
 @app.post("/generate-topics")
@@ -179,4 +262,5 @@ async def generate_topics(request: Request):
         'stream': True,
     }
 
-    return StreamingResponse(generate(json_data), media_type='text/event-stream')
+    selected_generator = random.choice([groqgenerate, generate])
+    return StreamingResponse(selected_generator(json_data), media_type='text/event-stream')