Spaces:
Sleeping
Sleeping
File size: 3,801 Bytes
f9e678f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import os
import streamlit as st
import requests
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from langchain_groq import ChatGroq
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import create_react_agent
from langchain_core.messages.ai import AIMessage
from fastapi.middleware.cors import CORSMiddleware
from threading import Thread
import uvicorn
# 🔹 Set API Keys
os.environ["GROQ_API_KEY"] = "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"
os.environ["TAVILY_API_KEY"] = "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"
# 🔹 Initialize FastAPI
app = FastAPI(title="LangGraph AI Agent")
# 🔹 Allow frontend to communicate with backend
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 🔹 Allowed model names
ALLOWED_MODEL_NAMES = ["llama3-70b-8192", "mixtral-8x7b-32768", "llama-3.3-70b-versatile", "gpt-4o-mini"]
# 🔹 Request Schema
class RequestState(BaseModel):
model_name: str
model_provider: str
system_prompt: str
messages: List[str]
allow_search: bool
# 🔹 AI Agent Function
def get_response_from_ai_agent(llm_id, query, allow_search, system_prompt, provider):
if provider == "Groq":
llm = ChatGroq(model=llm_id)
tools = [TavilySearchResults(max_results=5, api_key=os.environ["TAVILY_API_KEY"])] if allow_search else []
agent = create_react_agent(model=llm, tools=tools, state_modifier=system_prompt)
state = {"message": query}
response = agent.invoke(state)
messages = response.get("messages")
ai_messages = [message.content for message in messages if isinstance(message, AIMessage)]
return ai_messages[-1] if ai_messages else "No response from AI."
# 🔹 FastAPI Route
@app.post("/chat")
def chat_endpoint(request: RequestState):
if request.model_name not in ALLOWED_MODEL_NAMES:
return {"error": "Invalid model name"}
response = get_response_from_ai_agent(
llm_id=request.model_name,
query=request.messages,
allow_search=request.allow_search,
system_prompt=request.system_prompt,
provider=request.model_provider
)
return {"response": response}
# 🔹 Run FastAPI in a separate thread
def run_fastapi():
uvicorn.run(app, host="0.0.0.0", port=8000)
Thread(target=run_fastapi, daemon=True).start()
# 🔹 Streamlit UI
st.set_page_config(layout="wide", page_title="LangGraph AI Agent")
st.title("LangGraph AI Agent 🤖")
st.markdown("### Chat with an AI Model Powered by LangGraph!")
# 🔹 User Input Fields
system_prompt = st.text_area("Define your AI Agent:", placeholder="Type your prompt here...")
provider = st.selectbox("Select the model provider:", ["Groq", "OpenAI"])
model_name = st.selectbox(
"Select the model name:",
["llama3-70b-8192", "mixtral-8x7b-32768", "llama-3.3-70b-versatile"] if provider == "Groq" else ["gpt-4o-mini"]
)
allow_search = st.checkbox("Allow search")
user_query = st.text_area("Enter query:", placeholder="Type your query here...", height=150)
API_URL = "http://127.0.0.1:8000/chat"
if st.button("Chat 🤖"):
if user_query.strip():
response = requests.post(API_URL, json={
"model_name": model_name,
"model_provider": provider,
"system_prompt": system_prompt,
"messages": [user_query],
"allow_search": allow_search
})
if response.status_code == 200:
response_data = response.json()
if "error" in response_data:
st.error(response_data["error"])
else:
st.success(f"**AI Agent Response:** {response_data['response']}")
|