fastapi / main.py
ka1kuk's picture
Update main.py
dfbec24
raw
history blame
2.18 kB
import os
from langchain.memory import ConversationBufferMemory
from langchain.utilities import GoogleSearchAPIWrapper
from langchain.agents import AgentType, initialize_agent, Tool
from lang import G4F
from fastapi import FastAPI, Request
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from ImageCreator import generate_image_prodia
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"
),
]
llm = G4F(model=model)
agent_chain = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
@app.get("/")
def gello():
return "Hello! My name is Linlada."
@app.post('/linlada')
async def hello_post(request: Request):
llm = G4F(model=model)
data = await request.json()
prompt = data['prompt']
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
@app.route("/imagen", methods=["POST"])
def generate_image():
data = request.get_json()
prompt = data.get("prompt")
model = data.get("model")
sampler = data.get("sampler")
seed = int(data.get("seed"))
neg = data.get("neg")
response = generate_image_prodia(prompt, model, sampler, seed, neg)
return jsonify({"image": response})
@app.route("/test", methods=["POST"])
def test():
data = request.get_json()
prompt = data.get("prompt")
model = data.get("model")
sampler = data.get("sampler")
seed = int(data.get("seed"))
neg = data.get("neg")
res = [prompt, \n, model, \n, sampler, \n, seed, \n, neg]