1MR commited on
Commit
7cc5d4d
·
verified ·
1 Parent(s): b5c22a0

Create app.py

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