Update app.py

app.py

@@ -16,17 +16,21 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 
+
 # Set up basic configuration for logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Environment variables and configurations
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
+ACCOUNT_ID = os.environ.get("CLOUDFARE_ACCOUNT_ID")
+AUTH_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
 
 MODELS = [
     "mistralai/Mistral-7B-Instruct-v0.3",
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
-    "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+    "@cf/meta/llama-3.1-8b-instruct"
 ]
 
 # Initialize LlamaParse
@@ -81,33 +85,81 @@ def update_vectors(files, parser):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
+import os
+import json
+import requests
+from huggingface_hub import InferenceClient
+
+ACCOUNT_ID = "your-account-id"  # Replace with your actual Cloudflare account ID
+AUTH_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
+
 def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
-    client = InferenceClient(model, token=huggingface_token)
     full_response = ""
     messages = [{"role": "user", "content": prompt}]
     
-    for i in range(num_calls):
-        print(f"Starting API call {i+1}")
-        if should_stop:
-            print("Stop clicked, breaking loop")
-            break
-        try:
-            for message in client.chat_completion(
-                messages=messages,
-                max_tokens=max_tokens,
-                temperature=temperature,
-                stream=True,
-            ):
-                if should_stop:
-                    print("Stop clicked during streaming, breaking")
-                    break
-                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                    chunk = message.choices[0].delta.content
-                    full_response += chunk
-            print(f"API call {i+1} completed")
-        except Exception as e:
-            print(f"Error in generating response: {str(e)}")
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Cloudflare API
+        for i in range(num_calls):
+            print(f"Starting Cloudflare API call {i+1}")
+            if should_stop:
+                print("Stop clicked, breaking loop")
+                break
+            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": max_tokens,
+                        "temperature": temperature
+                    },
+                    stream=True
+                )
+                
+                for line in response.iter_lines():
+                    if should_stop:
+                        print("Stop clicked during streaming, breaking")
+                        break
+                    if line:
+                        try:
+                            json_data = json.loads(line.decode('utf-8').split('data: ')[1])
+                            chunk = json_data['response']
+                            full_response += chunk
+                        except json.JSONDecodeError:
+                            continue
+                print(f"Cloudflare API call {i+1} completed")
+            except Exception as e:
+                print(f"Error in generating response from Cloudflare: {str(e)}")
+    else:
+        # Original Hugging Face API logic
+        client = InferenceClient(model, token=huggingface_token)
+        
+        for i in range(num_calls):
+            print(f"Starting Hugging Face API call {i+1}")
+            if should_stop:
+                print("Stop clicked, breaking loop")
+                break
+            try:
+                for message in client.chat_completion(
+                    messages=messages,
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                    stream=True,
+                ):
+                    if should_stop:
+                        print("Stop clicked during streaming, breaking")
+                        break
+                    if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                        chunk = message.choices[0].delta.content
+                        full_response += chunk
+                print(f"Hugging Face API call {i+1} completed")
+            except Exception as e:
+                print(f"Error in generating response from Hugging Face: {str(e)}")
     
     # Clean up the response
     clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_response, flags=re.DOTALL)