1MB1 commited on
Commit
1ae7de2
·
verified ·
1 Parent(s): 4935255

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +177 -0
app.py ADDED
@@ -0,0 +1,177 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from llama_index.core import SimpleDirectoryReader
2
+ from llama_index.core.node_parser import SentenceSplitter
3
+ from llama_index.core import Settings
4
+ from llama_index.embeddings.huggingface import HuggingFaceEmbedding
5
+ from llama_index.core import VectorStoreIndex
6
+ from langchain_groq import ChatGroq
7
+ from langchain.tools import BaseTool, StructuredTool, tool
8
+ from pydantic import BaseModel
9
+ from langchain_community.tools.tavily_search import TavilySearchResults
10
+ from typing import TypedDict ,Annotated
11
+ from langchain_core.runnables import RunnablePassthrough
12
+ from langchain_core.output_parsers import StrOutputParser
13
+ import os
14
+ import uuid
15
+ from typing import TypedDict ,Annotated
16
+ from langchain_core.messages import AnyMessage,SystemMessage,HumanMessage,ToolMessage,AIMessage
17
+ import operator
18
+ from langgraph.checkpoint.memory import InMemorySaver
19
+ from langgraph.graph import StateGraph, END
20
+ from fastapi import FastAPI
21
+ import json
22
+ import shutil
23
+ import os
24
+ from fastapi import FastAPI, File, UploadFile
25
+ import time
26
+ import requests
27
+ from datetime import datetime
28
+ import os
29
+ os.environ["GROQ_API_KEY"] = "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"
30
+ os.environ["TAVILY_API_KEY"] = "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"
31
+
32
+ app = FastAPI()
33
+
34
+ @app.get("/")
35
+ def read_root():
36
+ app.state.vector_index = None
37
+ # keep_alive()
38
+ return {"message": "Connected"}
39
+
40
+ def keep_alive(space_url="https://1mr-apigmail.hf.space/ping", interval_hours=5):
41
+ while True:
42
+ try:
43
+ print(f"🔄 Pinging {space_url} at {datetime.now()}")
44
+ response = requests.get(space_url)
45
+
46
+ if response.status_code == 200:
47
+ print("")
48
+ else:
49
+ print("")
50
+ except Exception as e:
51
+ print("")
52
+
53
+ time.sleep(interval_hours * 3600)
54
+
55
+ # keep_alive()
56
+
57
+ @tool
58
+ def retrieve(query_text):
59
+ """
60
+ Retrieves relevant information from a vector index based on a query from reports.
61
+ Parameters:
62
+ - query_text (str): Query to search for relevant information.
63
+ Returns:
64
+ - str: Retrieved text from the document.
65
+ """
66
+ if not hasattr(app.state, "vector_index") or app.state.vector_index is None:
67
+ return "Vector index not found. Please upload a file first."
68
+ else:
69
+ retriever = app.state.vector_index.as_retriever(similarity_top_k=3)
70
+ result = retriever.retrieve(query_text)
71
+ if result:
72
+ return "\n\n".join([node.node.text for node in result])
73
+ return "No relevant information found."
74
+
75
+ tavily_search = TavilySearchResults(max_results=4)
76
+
77
+
78
+ @app.post("/uploadpdfs")
79
+ async def upload_file(file: UploadFile = File(...)):
80
+ # global vector_index
81
+
82
+ # Save uploaded file to a temp directory
83
+ temp_dir = "temp_uploads"
84
+ os.makedirs(temp_dir, exist_ok=True)
85
+ file_id = str(uuid.uuid4())
86
+ file_path = os.path.join(temp_dir, f"{file_id}_{file.filename}")
87
+
88
+ with open(file_path, "wb") as f:
89
+ shutil.copyfileobj(file.file, f)
90
+
91
+ # Load and parse document
92
+ documents = SimpleDirectoryReader(input_files=[file_path]).load_data()
93
+ parser = SentenceSplitter(chunk_size=300, chunk_overlap=50)
94
+ nodes = parser.get_nodes_from_documents(documents)
95
+
96
+ # Create or update vector index
97
+ embed_model = HuggingFaceEmbedding(model_name="WhereIsAI/UAE-Large-V1")
98
+ # if vector_index is None:
99
+ if not hasattr(app.state, "vector_index") or app.state.vector_index is None:
100
+ app.state.vector_index = VectorStoreIndex(nodes, embed_model=embed_model)
101
+ message = "New vector index created and file stored."
102
+ else:
103
+ app.state.vector_index.insert_nodes(nodes)
104
+ message = "File stored and vector index updated."
105
+
106
+ return {"message": message, "filename": file.filename}
107
+
108
+
109
+ class QueryRequest(BaseModel):
110
+ message: str
111
+
112
+ class AgentState(TypedDict):
113
+ messages: Annotated[list[AnyMessage], operator.add]
114
+
115
+ memory = InMemorySaver()
116
+
117
+ class Agent:
118
+ def __init__(self, model, tools, checkpointer=None, system=""):
119
+ self.system = system
120
+ graph = StateGraph(AgentState)
121
+ graph.add_node('llm',self.call_llm)
122
+ graph.add_node('action',self.take_action)
123
+ graph.add_conditional_edges("llm",self.exists_action,{True :"action",False:END})
124
+ graph.add_edge("action","llm")
125
+ graph.set_entry_point("llm")
126
+ self.graph = graph.compile(checkpointer=checkpointer)
127
+ self.tools = {t.name:t for t in tools}
128
+ self.model = model.bind_tools(tools)
129
+
130
+ def call_llm(self, state:AgentState):
131
+ messages = state['messages']
132
+ if self.system :
133
+ messages = [SystemMessage(content=self.system)] + messages
134
+ message = self.model.invoke(messages)
135
+ return {"messages":[message]}
136
+
137
+ def exists_action(self, state:AgentState):
138
+ result = state['messages'][-1]
139
+ return len(result.tool_calls) > 0
140
+
141
+ def take_action(self, state:AgentState):
142
+ tool_calls = state['messages'][-1].tool_calls
143
+ results = []
144
+ for t in tool_calls:
145
+ result= self.tools[t['name']].invoke(t['args'])
146
+ results.append(ToolMessage(tool_call_id=t['id'],name=t['name'],content=str(result)))
147
+ return {"messages":results}
148
+
149
+
150
+ system_Prompt="""
151
+ You are an AI assistant designed to assist users with health benefits, diet, nutrition information, and recipes.
152
+ You analyze patient reports to offer guidance on self-care with AI support.
153
+ Provide answers directly related to the question, without additional explanation or unrelated information.
154
+ """
155
+
156
+ tools=[retrieve,tavily_search]
157
+
158
+ model = ChatGroq(model="qwen-qwq-32b")
159
+
160
+ agent = Agent(model, tools, memory, system=system_Prompt)
161
+
162
+ thread = {"configurable": {'thread_id': '1'}}
163
+
164
+
165
+ @app.post("/askbot")
166
+ async def ask_question(query: QueryRequest):
167
+ messages = [HumanMessage(content=query.message)]
168
+ final_res = ""
169
+
170
+ for event in agent.graph.stream({'messages': messages}, thread):
171
+ for v in event.values():
172
+ if isinstance(v, dict) and 'messages' in v:
173
+ for msg in v['messages']:
174
+ if hasattr(msg, 'content') and isinstance(msg, AIMessage):
175
+ final_res += msg.content
176
+
177
+ return {"answer": final_res}