Update app.py

app.py

@@ -232,10 +232,18 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                 logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
-            for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
-                first_line = partial_response.split('\n')[0] if partial_response else ''
-                logging.info(f"Generated Response (first line): {first_line}")
-                yield partial_response
+            if model == "@cf/meta/llama-3.1-8b-instruct":
+                # Use Cloudflare API
+                for partial_response in get_response_from_cloudflare(message, num_calls=num_calls, temperature=temperature):
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+                    logging.info(f"Generated Response (first line): {first_line}")
+                    yield partial_response
+            else:
+                # Use Hugging Face API
+                for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+                    logging.info(f"Generated Response (first line): {first_line}")
+                    yield partial_response
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
         if "microsoft/Phi-3-mini-4k-instruct" in model:
@@ -245,6 +253,42 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
         else:
             yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
+def get_response_from_cloudflare(query, num_calls=3, temperature=0.2):
+    ACCOUNT_ID = "your-account-id"  # Replace with your actual Cloudflare account ID
+    AUTH_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
+    
+    prompt = f"Write a detailed and complete response that answers the following user question: '{query}'"
+    
+    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
+            )
+            
+            partial_response = ""
+            for line in response.iter_lines():
+                if line:
+                    try:
+                        json_data = json.loads(line.decode('utf-8').split('data: ')[1])
+                        chunk = json_data['response']
+                        partial_response += chunk
+                        yield partial_response
+                    except json.JSONDecodeError:
+                        continue
+        except Exception as e:
+            print(f"Error in generating response from Cloudflare: {str(e)}")
+
 def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
     context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n"