Bunnnyyy2005/Smart_engineering_RAG

1from fastapi import FastAPI, HTTPException
2from pydantic import BaseModel
3import uvicorn
4import os
5from dotenv import load_dotenv
6
7# Modern LangGraph and Stable LangChain Imports
8from langchain_groq import ChatGroq
9from langchain_core.tools import create_retriever_tool
10from langgraph.prebuilt import create_react_agent
11
12# Your Custom Modules
13from rag_engine import get_retriever
14from mcp_tools import get_live_machine_status
15
16# Load API Key
17load_dotenv()
18
19app = FastAPI(
20    title="Engineering Smart Agent API",
21    description="Bulletproof Backend for AI Troubleshooting Agent"
22)
23
24class QueryRequest(BaseModel):
25    query: str
26
27# 1. THE BRAIN: Using Mixtral - The most stable model for Tool Calling (No XML Bugs)
28llm = ChatGroq(
29    temperature=0, 
30    model_name="llama-3.1-8b-instant", 
31    groq_api_key=os.getenv("GROQ_API_KEY")
32)
33
34# 2. THE TOOLS: Connecting RAG and MCP
35retriever = get_retriever()
36rag_tool = create_retriever_tool(
37    retriever,
38    "engineering_manual_search",
39    "Searches technical manuals. Use this strictly when the user asks for theory, disadvantages, comparisons, or troubleshooting procedures."
40)
41
42tools = [rag_tool, get_live_machine_status]
43
44# 3. THE AGENT: Clean LangGraph Setup
45agent_executor = create_react_agent(llm, tools)
46
47# 4. SYSTEM PROMPT: Strict instructions
48SYSTEM_PROMPT = """You are a senior Engineering and AI Troubleshooting AI. 
49You have access to technical manuals and a live machine database. 
50- Use the live status tool ONLY if asked about a machine's current status. 
51- Use the manual search tool if asked about concepts, algorithms, disadvantages, or fixes.
52Always provide a clear, professional, and complete answer."""
53
54@app.get("/")
55def read_root():
56    return {"status": "Backend is running perfectly! 🚀"}
57
58@app.post("/ask")
59def ask_agent(request: QueryRequest):  # <-- Removed 'async' here!
60    try:
61        print(f"🚀 Received question: {request.query}")
62        print("🧠 Sending request to Groq API... (Please wait)")
63        
64        result = agent_executor.invoke({
65            "messages": [
66                ("system", SYSTEM_PROMPT),
67                ("user", request.query)
68            ]
69        })
70        
71        print("✅ Received response from Groq!")
72        final_answer = result["messages"][-1].content
73        
74        return {
75            "query": request.query,
76            "response": final_answer
77        }
78    except Exception as e:
79        print(f"❌ ERROR: {str(e)}")
80        raise HTTPException(status_code=500, detail=str(e))
81
82if __name__ == "__main__":
83    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)