Update main.py

main.py

@@ -1,7 +1,5 @@
 import os
-from langchain.memory import ConversationBufferMemory
-from langchain.utilities import GoogleSearchAPIWrapper
-from langchain.agents import AgentType, initialize_agent, Tool
+from googleapiclient.discovery import build
 from lang import G4F
 from fastapi import FastAPI, Request, Path
 from pydantic import BaseModel
@@ -23,38 +21,64 @@ google_api_key = os.environ["GOOGLE_API_KEY"]
 cse_id = os.environ["GOOGLE_CSE_ID"]
 model = os.environ['default_model']
 
-
-search = GoogleSearchAPIWrapper()
-tools = [
-    Tool(
-    name ="Search" ,
-    func=search.run,
-    description="useful when you need to answer questions about current events"
-    ),
-]
-
-llm = G4F(model=model)
-agent_chain = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
+def search_google(query):
+    service = build("customsearch", "v1", developerKey=google_api_key)
+    result = service.cse().list(
+        q=query,
+        cx=cse_id
+    ).execute()
+    return result['items']
 
 
 @app.get("/")
 def hello():
     return "Hello! My name is Linlada."
 
+
+
+def gpt_with_google_search(prompt):
+    search_results = search_google(prompt)
+    text = ""
+    ref = ""
+    for item in search_results:
+        text += item['title'] + "\n" + item['snippet'] + "\n\n"
+        ref += "- {} ({})\n".format(item['title'], item['link'])
+    results = generate_text(f'Summarize: {text}')
+    res = "{} \n\n {}".format(results, ref)
+    return res
+
+def gpt_with_google_search(prompt):
+    search_results = search_google(prompt)
+    text = ""
+    ref = ""
+    ref_count = 1
+    for item in search_results:
+        text += item['title'] + "\n" + item['snippet'] + "\n\n"
+        ref += " [{}] {} ({})\n".format(ref_count, item['title'], item['link'])
+        ref_count += 1
+    results = generate_text(f'Summarize: {text}')
+    res = "{} \n\n{}".format(results, ref)
+    return res
+
+class Linlada(BaseModel):
+    prompt: str
+    web_access: str = False
+    sampler: str = None
+    seed: int = None
+    neg: str = None
+
+
 @app.post('/linlada')
-async def hello_post(request: Request):
+def linlada(request: Linlada):
     llm = G4F(model=model)
-    data = await request.json()
-    prompt = data['prompt']
-    chat = llm(prompt)
+    prompt = request.prompt
+    web_access = request.web_access
+    if web_access == True:
+       chat = gpt_with_google_search(prompt)
+    else:
+        chat = llm(prompt)
     return chat
 
-@app.post('/search')
-async def searches(request: Request):
-    data = await request.json()
-    prompt = data['prompt']
-    response = agent_chain.run(input=prompt)
-    return response
 
 class User(BaseModel):
     prompt: str