Cheh Kit Hong commited on
Commit
94e0eef
Β·
1 Parent(s): 5121ff0

changed rag method flags

Browse files
Files changed (6) hide show
  1. README.md +5 -5
  2. agent/graph.py +2 -2
  3. agent/nodes.py +1 -1
  4. agent/prompts.py +18 -7
  5. agent/state.py +0 -10
  6. config.py +1 -1
README.md CHANGED
@@ -151,8 +151,8 @@ nano .env # or use your preferred editor
151
 
152
  **Required environment variables:**
153
  ```bash
154
- ANTHROPIC_API_KEY=sk-ant-xxxxxxxxxxxxx
155
- TAVILY_API_KEY=tvly-xxxxxxxxxxxxx
156
  ```
157
 
158
  ### 5️⃣ Prepare Data
@@ -171,7 +171,6 @@ mkdir -p data/documents data/chroma_db
171
  python main.py
172
  ```
173
 
174
- The app will be available at: **http://127.0.0.1:7860**
175
 
176
  ### 7️⃣ Run Tests (Optional)
177
  ```bash
@@ -283,11 +282,12 @@ lsof -ti:7860 | xargs kill -9
283
 
284
  - [UV Documentation](https://github.com/astral-sh/uv)
285
  - [LangGraph Docs](https://langchain-ai.github.io/langgraph/)
286
- - [Anthropic API](https://docs.anthropic.com/)
287
  - [Tavily API](https://docs.tavily.com/)
288
  - [ChromaDB Docs](https://docs.trychroma.com/)
289
 
290
- ## πŸ“š Reference
 
291
  1. Wei, H., Sun, Y., & Li, Y. (2025). Deepseek-ocr: Contexts optical compression. arXiv preprint arXiv:2510.18234.
292
  2. Chen, X., Chu, F. J., Gleize, P., Liang, K. J., Sax, A., Tang, H., ... & SAM 3D Team. (2025). SAM 3D: 3Dfy Anything in Images. arXiv preprint arXiv:2511.16624.
293
  3. Carion, N., Gustafson, L., Hu, Y. T., Debnath, S., Hu, R., Suris, D., ... & Feichtenhofer, C. (2025). SAM 3: Segment Anything with Concepts. arXiv preprint arXiv:2511.16719.
 
151
 
152
  **Required environment variables:**
153
  ```bash
154
+ GOOGLE_API_KEY=xxxxxxxxxxxxx # Gemini API key
155
+ TAVILY_API_KEY=tvly-xxxxxxxxxxxxx # Enable web search
156
  ```
157
 
158
  ### 5️⃣ Prepare Data
 
171
  python main.py
172
  ```
173
 
 
174
 
175
  ### 7️⃣ Run Tests (Optional)
176
  ```bash
 
282
 
283
  - [UV Documentation](https://github.com/astral-sh/uv)
284
  - [LangGraph Docs](https://langchain-ai.github.io/langgraph/)
285
+ - [Gemini API](https://ai.google.dev/gemini-api/docs/api-key)
286
  - [Tavily API](https://docs.tavily.com/)
287
  - [ChromaDB Docs](https://docs.trychroma.com/)
288
 
289
+ ## πŸ“š Reference (Used for demo as document in vector store)
290
+
291
  1. Wei, H., Sun, Y., & Li, Y. (2025). Deepseek-ocr: Contexts optical compression. arXiv preprint arXiv:2510.18234.
292
  2. Chen, X., Chu, F. J., Gleize, P., Liang, K. J., Sax, A., Tang, H., ... & SAM 3D Team. (2025). SAM 3D: 3Dfy Anything in Images. arXiv preprint arXiv:2511.16624.
293
  3. Carion, N., Gustafson, L., Hu, Y. T., Debnath, S., Hu, R., Suris, D., ... & Feichtenhofer, C. (2025). SAM 3: Segment Anything with Concepts. arXiv preprint arXiv:2511.16719.
agent/graph.py CHANGED
@@ -40,8 +40,8 @@ def create_agent_graph(llm, vectordb, search_tools) -> StateGraph:
40
  "web_search_agent_node",
41
  tools_condition,
42
  {
43
- "tools": "web_search_tool_node", # Changed key from node name to "tools"
44
- "__end__": "generate_node", # Changed key from "generate_node" to "__end__"
45
  }
46
  )
47
 
 
40
  "web_search_agent_node",
41
  tools_condition,
42
  {
43
+ "tools": "web_search_tool_node",
44
+ "__end__": "generate_node",
45
  }
46
  )
47
 
agent/nodes.py CHANGED
@@ -20,7 +20,7 @@ def router_node(state: AgentState, llm):
20
 
21
  def routing_logic(state: AgentState) -> str:
22
  rag_method = state["rag_method"]
23
- if rag_method == "RAG":
24
  return "vectordb_node"
25
  elif rag_method == "WEBSEARCH":
26
  return "web_search_agent_node"
 
20
 
21
  def routing_logic(state: AgentState) -> str:
22
  rag_method = state["rag_method"]
23
+ if rag_method == "VECTORDB":
24
  return "vectordb_node"
25
  elif rag_method == "WEBSEARCH":
26
  return "web_search_agent_node"
agent/prompts.py CHANGED
@@ -15,18 +15,29 @@ def determine_rag_method_prompt() -> str:
15
  You are a query classification model. Given the user's query, you must classify the method to use
16
  as one and only one of the following options:
17
 
18
- 1. **RAG**: The query asks about specific documents, papers, or systems like DeepAnalyze, AgentMem, SAM3, SAM 3, SAM3D, DeepSeek-OCR, or any technical architecture/implementation details from research papers.
19
- 2. **WEBSEARCH**: The query asks for current events, latest news, real-time information after January 2024, or general factual knowledge not in specialized documents.
20
- 3. **GENERAL**: The query is a simple calculation, definition, reasoning task, or common knowledge question that doesn't need external data.
 
 
 
 
 
 
21
 
22
  **Examples:**
23
- - "What is DeepAnalyze?" β†’ RAG
24
- - "Explain SAM 3 architecture" β†’ RAG
25
  - "Latest AI news in 2025" β†’ WEBSEARCH
26
  - "What is 15 times 7?" β†’ GENERAL
27
 
28
- Respond STRICTLY with only one word: RAG, WEBSEARCH, or GENERAL. No punctuation or extra text.
29
  """
30
-
 
 
 
 
 
31
  if __name__ == "__main__":
32
  pass
 
15
  You are a query classification model. Given the user's query, you must classify the method to use
16
  as one and only one of the following options:
17
 
18
+ 1. **VECTORDB**:
19
+ The query asks about specific documents, papers, or systems.
20
+ This includes technical architecture/implementation details from research papers for DeepAnalyze, AgentMem, SAM3, SAM 3, SAM3D, DeepSeek-OCR
21
+
22
+ 2. **WEBSEARCH**:
23
+ The query asks for current events, latest news, real-time information after January 2024, or general factual knowledge not in specialized documents.
24
+
25
+ 3. **GENERAL**:
26
+ The query is a simple calculation, definition, reasoning task, or common knowledge question that doesn't need external data.
27
 
28
  **Examples:**
29
+ - "What is DeepAnalyze?" β†’ VECTORDB
30
+ - "Explain SAM 3 architecture" β†’ VECTORDB
31
  - "Latest AI news in 2025" β†’ WEBSEARCH
32
  - "What is 15 times 7?" β†’ GENERAL
33
 
34
+ Respond STRICTLY with only one word: VECTORDB, WEBSEARCH, or GENERAL. No punctuation or extra text.
35
  """
36
+ # For dynamic RAG routing based on user query, each route decision need to clafify what
37
+ # information can be retrieved from which source.
38
+ # If more than 2 vectorstores are used, the routing logic can to be updated accordingly.
39
+ # Example: VECTORDB1, VECTORDB2, WEBSEARCH, GENERAL
40
+
41
+
42
  if __name__ == "__main__":
43
  pass
agent/state.py CHANGED
@@ -14,16 +14,6 @@ class AgentState(TypedDict):
14
  rag_method: RAG_method
15
  context: Optional[str]
16
 
17
- questionIsClear: bool
18
- conversation_summary: str = ""
19
-
20
-
21
- # Implement later if needed, omit first
22
- class QueryAnalysis(BaseModel):
23
- """Structured output for query analysis"""
24
- is_clear: bool = Field(description="Indicates if the user's question is clear and answerable")
25
- questions: List[str] = Field(description="List of rewritten, self-contained questions")
26
- clarification_needed: str = Field(description="Explanation if the question is unclear")
27
 
28
  if __name__ == "__main__":
29
  pass
 
14
  rag_method: RAG_method
15
  context: Optional[str]
16
 
 
 
 
 
 
 
 
 
 
 
17
 
18
  if __name__ == "__main__":
19
  pass
config.py CHANGED
@@ -2,7 +2,7 @@ import os
2
 
3
  configs = {
4
  "DATA_PATH": "./docs/markdowns",
5
- "PERSIST_PATH": "./chroma_data",
6
  "EMBEDDING_MODEL_NAME": "sentence-transformers/all-mpnet-base-v2",
7
  "COLLECTION_NAME": "langchain_mpnet_collection",
8
  "LLM_MODEL_NAME": "gemini-2.0-flash",
 
2
 
3
  configs = {
4
  "DATA_PATH": "./docs/markdowns",
5
+ "PERSIST_PATH": "./chroma_data", # modify if want to create another vectorstore like chroma_data_2
6
  "EMBEDDING_MODEL_NAME": "sentence-transformers/all-mpnet-base-v2",
7
  "COLLECTION_NAME": "langchain_mpnet_collection",
8
  "LLM_MODEL_NAME": "gemini-2.0-flash",