EtienneB commited on
Commit
64434a5
·
1 Parent(s): 600dd01
Files changed (3) hide show
  1. agent.py +89 -37
  2. app.py +0 -3
  3. requirements.txt +14 -9
agent.py CHANGED
@@ -2,8 +2,7 @@ import os
2
 
3
  from dotenv import load_dotenv
4
  from langchain_community.vectorstores import Chroma
5
- from langchain_core.messages import HumanMessage, SystemMessage
6
- from langchain_core.tools import tool
7
  from langchain_huggingface import (ChatHuggingFace, HuggingFaceEmbeddings,
8
  HuggingFaceEndpoint)
9
  from langgraph.graph import START, MessagesState, StateGraph
@@ -33,7 +32,6 @@ tools = [
33
  wiki_search, analyze_excel_file, arvix_search, audio_transcription, python_code_parser
34
  ]
35
 
36
-
37
  # Load system prompt
38
  system_prompt = """
39
  You are a helpful assistant tasked with answering questions using a set of tools.
@@ -46,62 +44,106 @@ Your answer should only start with "FINAL ANSWER: ", then follows with the answe
46
  # System message
47
  sys_msg = SystemMessage(content=system_prompt)
48
 
49
- # Embeddings + Chroma Vector Store
50
- embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
51
- vector_store = Chroma(
52
- collection_name="langgraph-documents",
53
- embedding_function=embeddings,
54
- persist_directory="chroma_db" # Use a persistent directory
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55
  )
56
 
 
57
  def build_graph():
58
- """Build the graph"""
59
 
60
  # First create the HuggingFaceEndpoint
61
  llm_endpoint = HuggingFaceEndpoint(
62
- # repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
63
- repo_id="mistralai/Mistral-7B-Instruct-v0.2",
64
- # Other models to try:
65
- # "meta-llama/Llama-2-7b-chat-hf"
66
- # "google/gemma-7b-it"
67
- # "mosaicml/mpt-7b-instruct"
68
- # "tiiuae/falcon-7b-instruct"
69
  huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
70
- temperature=0.1, # Lower temperature for more consistent responses
71
- max_new_tokens=1024,
72
- timeout=30,
73
- )
74
-
75
  # Then wrap it with ChatHuggingFace to get chat model functionality
76
  llm = ChatHuggingFace(llm=llm_endpoint)
77
 
78
  # Bind tools to LLM
79
  llm_with_tools = llm.bind_tools(tools)
80
 
81
- # Node
82
  def assistant(state: MessagesState):
83
  """Assistant node"""
84
- return {"messages": [llm_with_tools.invoke([system_prompt] + state["messages"])]}
85
-
86
- def retriever(state: MessagesState):
87
- similar = vector_store.similarity_search(state["messages"][0].content)
88
- if similar:
89
- example_msg = HumanMessage(content=f"Here is a similar question:\n\n{similar[0].page_content}")
90
- return {"messages": [sys_msg] + state["messages"] + [example_msg]}
91
- return {"messages": [sys_msg] + state["messages"]}
 
 
 
 
 
 
 
 
 
 
 
 
 
92
 
 
93
  builder = StateGraph(MessagesState)
94
- builder.add_node("retriever", retriever)
95
  builder.add_node("assistant", assistant)
96
  builder.add_node("tools", ToolNode(tools))
 
97
  builder.add_edge(START, "retriever")
98
  builder.add_edge("retriever", "assistant")
99
  builder.add_conditional_edges("assistant", tools_condition)
100
  builder.add_edge("tools", "assistant")
101
-
102
- # Compile graph
103
  return builder.compile()
104
 
 
105
  # test
106
  if __name__ == "__main__":
107
  question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
@@ -109,6 +151,16 @@ if __name__ == "__main__":
109
  graph = build_graph()
110
  # Run the graph
111
  messages = [HumanMessage(content=question)]
112
- messages = graph.invoke({"messages": messages})
113
- for m in messages["messages"]:
114
- m.pretty_print()
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  from dotenv import load_dotenv
4
  from langchain_community.vectorstores import Chroma
5
+ from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
 
6
  from langchain_huggingface import (ChatHuggingFace, HuggingFaceEmbeddings,
7
  HuggingFaceEndpoint)
8
  from langgraph.graph import START, MessagesState, StateGraph
 
32
  wiki_search, analyze_excel_file, arvix_search, audio_transcription, python_code_parser
33
  ]
34
 
 
35
  # Load system prompt
36
  system_prompt = """
37
  You are a helpful assistant tasked with answering questions using a set of tools.
 
44
  # System message
45
  sys_msg = SystemMessage(content=system_prompt)
46
 
47
+
48
+ def get_vector_store(persist_directory="chroma_db"):
49
+ """
50
+ Initializes and returns a Chroma vector store.
51
+ If the database exists, it loads it. If not, it creates it,
52
+ adds some initial documents, and persists them.
53
+ """
54
+ embedding_function = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
55
+
56
+ if os.path.exists(persist_directory) and os.listdir(persist_directory):
57
+ print("Loading existing vector store...")
58
+ vector_store = Chroma(
59
+ persist_directory=persist_directory,
60
+ embedding_function=embedding_function
61
+ )
62
+ else:
63
+ print("Creating new vector store...")
64
+ os.makedirs(persist_directory, exist_ok=True)
65
+ # Example documents to add
66
+ initial_documents = [
67
+ "The Principle of Double Effect is an ethical theory that distinguishes between the intended and foreseen consequences of an action.",
68
+ "St. Thomas Aquinas is often associated with the development of the Principle of Double Effect.",
69
+ "LangGraph is a library for building stateful, multi-actor applications with LLMs.",
70
+ "Chroma is a vector database used for storing and retrieving embeddings."
71
+ ]
72
+ vector_store = Chroma.from_texts(
73
+ texts=initial_documents,
74
+ embedding=embedding_function,
75
+ persist_directory=persist_directory
76
+ )
77
+ # No need to call persist() when using from_texts with a persist_directory
78
+
79
+ return vector_store
80
+
81
+ # --- Initialize Vector Store and Retriever ---
82
+ vector_store = get_vector_store()
83
+ retriever_component = vector_store.as_retriever(
84
+ search_type="mmr", # Use Maximum Marginal Relevance for diverse results
85
+ search_kwargs={'k': 2, 'lambda_mult': 0.5} # Retrieve 2 documents
86
  )
87
 
88
+
89
  def build_graph():
90
+ """Build the graph"""
91
 
92
  # First create the HuggingFaceEndpoint
93
  llm_endpoint = HuggingFaceEndpoint(
94
+ repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
 
 
 
 
 
 
95
  huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
96
+ temperature=0.3,
97
+ max_new_tokens=2048,
98
+ timeout=60,
99
+ )
100
+
101
  # Then wrap it with ChatHuggingFace to get chat model functionality
102
  llm = ChatHuggingFace(llm=llm_endpoint)
103
 
104
  # Bind tools to LLM
105
  llm_with_tools = llm.bind_tools(tools)
106
 
107
+ # --- Nodes ---
108
  def assistant(state: MessagesState):
109
  """Assistant node"""
110
+ # Prepend the system message to the state
111
+ messages_with_system_prompt = [sys_msg] + state["messages"]
112
+ return {"messages": [llm_with_tools.invoke(messages_with_system_prompt)]}
113
+
114
+ def retriever_node(state: MessagesState):
115
+ """
116
+ Retrieves relevant documents from the vector store based on the latest human message.
117
+ """
118
+ last_human_message = state["messages"][-1].content
119
+ retrieved_docs = retriever_component.invoke(last_human_message)
120
+
121
+ if retrieved_docs:
122
+ retrieved_context = "\n\n".join([doc.page_content for doc in retrieved_docs])
123
+ # Create a ToolMessage to hold the retrieved context
124
+ context_message = ToolMessage(
125
+ content=f"Retrieved context from vector store:\n\n{retrieved_context}",
126
+ tool_call_id="retriever" # A descriptive ID
127
+ )
128
+ return {"messages": [context_message]}
129
+
130
+ return {"messages": []}
131
 
132
+ # --- Graph Definition ---
133
  builder = StateGraph(MessagesState)
134
+ builder.add_node("retriever", retriever_node)
135
  builder.add_node("assistant", assistant)
136
  builder.add_node("tools", ToolNode(tools))
137
+
138
  builder.add_edge(START, "retriever")
139
  builder.add_edge("retriever", "assistant")
140
  builder.add_conditional_edges("assistant", tools_condition)
141
  builder.add_edge("tools", "assistant")
142
+
143
+ # Compile graph
144
  return builder.compile()
145
 
146
+
147
  # test
148
  if __name__ == "__main__":
149
  question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
 
151
  graph = build_graph()
152
  # Run the graph
153
  messages = [HumanMessage(content=question)]
154
+ # The initial state for the graph
155
+ initial_state = {"messages": messages}
156
+
157
+ # Invoke the graph stream to see the steps
158
+ for s in graph.stream(initial_state, stream_mode="values"):
159
+ message = s["messages"][-1]
160
+ if isinstance(message, ToolMessage):
161
+ print("---RETRIEVED CONTEXT---")
162
+ print(message.content)
163
+ print("-----------------------")
164
+ else:
165
+ message.pretty_print()
166
+
app.py CHANGED
@@ -36,9 +36,6 @@ class BasicAgent:
36
 
37
 
38
 
39
-
40
-
41
-
42
  def run_and_submit_all( profile: gr.OAuthProfile | None):
43
  """
44
  Fetches all questions, runs the BasicAgent on them, submits all answers,
 
36
 
37
 
38
 
 
 
 
39
  def run_and_submit_all( profile: gr.OAuthProfile | None):
40
  """
41
  Fetches all questions, runs the BasicAgent on them, submits all answers,
requirements.txt CHANGED
@@ -2,6 +2,7 @@
2
  gradio
3
  requests
4
  pandas
 
5
 
6
  # LangChain and ecosystem
7
  langchain
@@ -10,24 +11,28 @@ langchain-community
10
  langgraph
11
  langchain-huggingface
12
  langchain-chroma
 
13
  sentence-transformers
14
 
15
  # Hugging Face integration
16
  huggingface_hub
17
  transformers
18
- accelerate # Needed for many transformer-based models
19
 
20
  # Environment config
21
  python-dotenv
22
 
23
  # Tools dependencies
24
- duckduckgo-search # Required for web_search tool
25
- pytz # Required for get_current_time_in_timezone tool
 
 
 
 
 
26
 
27
- # Additional utilities for better error handling and performance
28
  typing-extensions
29
- asyncio-throttle # For rate limiting (optional)
30
- tenacity # For retry logic (optional)
31
-
32
- # Optional: For better logging and monitoring
33
- loguru # Better logging (optional)
 
2
  gradio
3
  requests
4
  pandas
5
+ openpyxl # For reading excel files with pandas
6
 
7
  # LangChain and ecosystem
8
  langchain
 
11
  langgraph
12
  langchain-huggingface
13
  langchain-chroma
14
+ chromadb # Explicitly add the Chroma database
15
  sentence-transformers
16
 
17
  # Hugging Face integration
18
  huggingface_hub
19
  transformers
20
+ accelerate
21
 
22
  # Environment config
23
  python-dotenv
24
 
25
  # Tools dependencies
26
+ duckduckgo-search
27
+ pytz
28
+ wikipedia # For WikipediaLoader
29
+ arxiv # For ArxivLoader
30
+ assemblyai # For AssemblyAIAudioTranscriptLoader
31
+ tree-sitter # For LanguageParser
32
+ tree-sitter-languages # For LanguageParser
33
 
34
+ # Additional utilities
35
  typing-extensions
36
+ asyncio-throttle
37
+ tenacity
38
+ loguru