Spaces:
Sleeping
Sleeping
| 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 | |
| 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']}") | |