Update app.py

app.py

@@ -298,21 +298,57 @@ def get_response_with_search(query, model, num_calls=3, temperature=0.2):
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
 After writing the document, please provide a list of sources used in your response."""
-    
-    client = InferenceClient(model, token=huggingface_token)
-    
-    main_content = ""
-    for i in range(num_calls):
-        for message in client.chat_completion(
-            messages=[{"role": "user", "content": prompt}],
-            max_tokens=1000,
-            temperature=temperature,
-            stream=True,
-        ):
-            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                chunk = message.choices[0].delta.content
-                main_content += chunk
-                yield main_content, ""  # Yield partial main content without sources
+
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Use Cloudflare API
+        ACCOUNT_ID = "your-account-id"  # Replace with your actual Cloudflare account ID
+        AUTH_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
+        
+        main_content = ""
+        for i in range(num_calls):
+            try:
+                response = requests.post(
+                    f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/meta/llama-3.1-8b-instruct",
+                    headers={"Authorization": f"Bearer {AUTH_TOKEN}"},
+                    json={
+                        "stream": True,
+                        "messages": [
+                            {"role": "system", "content": "You are a friendly assistant"},
+                            {"role": "user", "content": prompt}
+                        ],
+                        "max_tokens": 1000,
+                        "temperature": temperature
+                    },
+                    stream=True
+                )
+                
+                for line in response.iter_lines():
+                    if line:
+                        try:
+                            json_data = json.loads(line.decode('utf-8').split('data: ')[1])
+                            chunk = json_data['response']
+                            main_content += chunk
+                            yield main_content, ""  # Yield partial main content without sources
+                        except json.JSONDecodeError:
+                            continue
+            except Exception as e:
+                print(f"Error in generating response from Cloudflare: {str(e)}")
+    else:
+        # Use Hugging Face API
+        client = InferenceClient(model, token=huggingface_token)
+        
+        main_content = ""
+        for i in range(num_calls):
+            for message in client.chat_completion(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=1000,
+                temperature=temperature,
+                stream=True,
+            ):
+                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                    chunk = message.choices[0].delta.content
+                    main_content += chunk
+                    yield main_content, ""  # Yield partial main content without sources
 
 def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
     embed = get_embeddings()