import os
from langchain.memory import ConversationBufferMemory
from langchain.utilities import GoogleSearchAPIWrapper
from langchain.agents import initialize_agent, Tool
from lang import G4F
from fastapi import FastAPI
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware

app = FastAPI()

app.add_middleware(  # add the middleware
    CORSMiddleware,
    allow_credentials=True,  # allow credentials
    allow_origins=["*"],  # allow all origins
    allow_methods=["*"],  # allow all methods
    allow_headers=["*"],  # allow all headers
)


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"
    ),
]

memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

llm = G4F(model=model)
agent_chain = initialize_agent(tools, llm, agent="chat-conversational-react-description",
                               verbose=True, memory=memory)

class messages(BaseModel):
    messages: str

@app.get("/")
def gello():
    return "Hello! My name is Linlada."

@app.post('/linlada')
def hello_post(message: messages):
    llm = G4F(model=model)
    chat = llm(message)
    return chat

@app.post('/search')
def searches(message: messages):
    response = agent_chain.run(input=message)
    return response