Update main.py

main.py

@@ -39,14 +39,20 @@ async def favicon():
     favicon_path = Path(__file__).parent / "favicon.ico"
     return FileResponse(favicon_path, media_type="image/x-icon")
     
-def generate_search(query: str, stream: bool = True) -> str:
+def generate_search(query: str, systemprompt: Optional[str] = None, stream: bool = True) -> str:
     headers = {"User-Agent": ""}
+    
+    # Use the provided system prompt, or default to "Be Helpful and Friendly"
+    system_message = systemprompt or "Be Helpful and Friendly"
+    
+    # Create the prompt history with the user query and system message
     prompt = [
         {"role": "user", "content": query},
     ]
     
-    prompt.insert(0, {"content": "Be Helpful and Friendly", "role": "system"})
+    prompt.insert(0, {"content": system_message, "role": "system"})
     
+    # Prepare the payload for the API request
     payload = {
         "is_vscode_extension": True,
         "message_history": prompt,
@@ -54,10 +60,12 @@ def generate_search(query: str, stream: bool = True) -> str:
         "user_input": prompt[-1]["content"],
     }
     
-    chat_endpoint = secret_api_endpoint_3
-    response = requests.post(chat_endpoint, headers=headers, json=payload, stream=True)
+    # Send the request to the chat endpoint
+    response = requests.post(secret_api_endpoint_3, headers=headers, json=payload, stream=True)
     
     streaming_text = ""
+    
+    # Process the streaming response
     for value in response.iter_lines(decode_unicode=True):
         if value.startswith("data: "):  
             try:
@@ -89,21 +97,21 @@ def generate_search(query: str, stream: bool = True) -> str:
     if not stream:
         yield streaming_text
 
+
 @app.get("/searchgpt")
-async def search_gpt(q: str, stream: Optional[bool] = False):
+async def search_gpt(q: str, stream: Optional[bool] = False, systemprompt: Optional[str] = None):
     if not q:
         raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
     
     if stream:
         return StreamingResponse(
-            generate_search(q, stream=True),
+            generate_search(q, systemprompt=systemprompt, stream=True),
             media_type="text/event-stream"
         )
     else:
         # For non-streaming, collect the text and return as JSON response
-        response_text = "".join([chunk for chunk in generate_search(q, stream=False)])
+        response_text = "".join([chunk for chunk in generate_search(q, systemprompt=systemprompt, 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)