Update main.py

main.py

@@ -1,9 +1,13 @@
 import os
 from dotenv import load_dotenv
 from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
 from pydantic import BaseModel
 import httpx
+import requests
+import re
+import json
+from typing import Optional
 
 load_dotenv()
 
@@ -14,9 +18,10 @@ api_keys_str = os.getenv('API_KEYS')
 valid_api_keys = api_keys_str.split(',') if api_keys_str else []
 secret_api_endpoint = os.getenv('SECRET_API_ENDPOINT')
 secret_api_endpoint_2 = os.getenv('SECRET_API_ENDPOINT_2')
+secret_api_endpoint_3 = os.getenv('SECRET_API_ENDPOINT_3')  # New endpoint for searchgpt
 
-# Validate if the main secret API endpoint is set
-if not secret_api_endpoint or not secret_api_endpoint_2:
+# Validate if the main secret API endpoints are set
+if not secret_api_endpoint or not secret_api_endpoint_2 or not secret_api_endpoint_3:
     raise HTTPException(status_code=500, detail="API endpoint(s) are not configured in environment variables.")
 
 # Define models that should use the secondary endpoint
@@ -27,6 +32,58 @@ class Payload(BaseModel):
     messages: list
     stream: bool
 
+def generate_search(query: str, stream: bool = True) -> str:
+    headers = {"User-Agent": ""}
+    prompt = [
+        {"role": "user", "content": query},
+    ]
+    
+    # Insert the system prompt at the beginning of the conversation history
+    prompt.insert(0, {"content": "Be Helpful and Friendly coz ur searchgpt", "role": "system"})
+    
+    payload = {
+        "is_vscode_extension": True,
+        "message_history": prompt,
+        "requested_model": "searchgpt",
+        "user_input": prompt[-1]["content"],
+    }
+    
+    # Use the newly added SECRET_API_ENDPOINT_3 for the search API call
+    chat_endpoint = secret_api_endpoint_3
+    response = requests.post(chat_endpoint, headers=headers, json=payload, stream=True)
+    
+    # Collect streamed text content
+    streaming_text = ""
+    for value in response.iter_lines(decode_unicode=True, chunk_size=12):
+        modified_value = re.sub("data:", "", value)
+        if modified_value:
+            try:
+                json_modified_value = json.loads(modified_value)
+                content = json_modified_value["choices"][0]["delta"]["content"]
+                if stream:
+                    yield f"data: {content}\n\n"
+                streaming_text += content
+            except:
+                continue
+    
+    if not stream:
+        yield streaming_text
+
+@app.get("/searchgpt")
+async def search_gpt(q: str, stream: Optional[bool] = True):
+    if not q:
+        raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
+    
+    if stream:
+        return StreamingResponse(
+            generate_search(q, stream=True),
+            media_type="text/event-stream"
+        )
+    else:
+        # For non-streaming response, collect all content and return as JSON
+        response_text = "".join([chunk for chunk in generate_search(q, stream=False)])
+        return JSONResponse(content={"response": response_text})
+
 @app.get("/", response_class=HTMLResponse)
 async def root():
     # Open and read the content of index.html (in the same folder as the app)
@@ -54,14 +111,7 @@ async def fetch_models():
 
 @app.post("/chat/completions")
 async def get_completion(payload: Payload, request: Request):
-    api_key = request.headers.get("Authorization")
-    print(f"API Key Used: {api_key}")
-    print(f"Raw Payload: {payload.dict()}")
-    # Validate API key
-    if api_key not in valid_api_keys:
-        raise HTTPException(status_code=403, detail="Forbidden: Invalid API key. Join dsc.gg/chadgang and DM @mr_leaderyt on Discord for a free API key :)")
-    
-    # Determine which endpoint to use based on the model
+    # Use the correct endpoint depending on the model type (no authentication now 😉)
     endpoint = secret_api_endpoint_2 if payload.model in alternate_models else secret_api_endpoint
 
     # Use the payload directly as it includes stream and other user data
@@ -89,9 +139,9 @@ async def startup_event():
     print("API endpoints:")
     print("GET /")
     print("GET /models")
-    print("POST /v1/chat/completions")
+    print("GET /searchgpt")  # We now have the new search API
+    print("POST /chat/completions")
 
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)
-