flowcept 0.8.11__py3-none-any.whl → 0.8.12__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- flowcept/__init__.py +7 -4
- flowcept/agents/__init__.py +5 -0
- flowcept/{flowceptor/consumers/agent/client_agent.py → agents/agent_client.py} +22 -12
- flowcept/agents/agents_utils.py +181 -0
- flowcept/agents/dynamic_schema_tracker.py +191 -0
- flowcept/agents/flowcept_agent.py +30 -0
- flowcept/agents/flowcept_ctx_manager.py +175 -0
- flowcept/agents/gui/__init__.py +5 -0
- flowcept/agents/gui/agent_gui.py +76 -0
- flowcept/agents/gui/gui_utils.py +239 -0
- flowcept/agents/llms/__init__.py +1 -0
- flowcept/agents/llms/claude_gcp.py +139 -0
- flowcept/agents/llms/gemini25.py +119 -0
- flowcept/agents/prompts/__init__.py +1 -0
- flowcept/{flowceptor/adapters/agents/prompts.py → agents/prompts/general_prompts.py} +18 -0
- flowcept/agents/prompts/in_memory_query_prompts.py +297 -0
- flowcept/agents/tools/__init__.py +1 -0
- flowcept/agents/tools/general_tools.py +102 -0
- flowcept/agents/tools/in_memory_queries/__init__.py +1 -0
- flowcept/agents/tools/in_memory_queries/in_memory_queries_tools.py +704 -0
- flowcept/agents/tools/in_memory_queries/pandas_agent_utils.py +309 -0
- flowcept/cli.py +286 -44
- flowcept/commons/daos/docdb_dao/mongodb_dao.py +47 -0
- flowcept/commons/daos/mq_dao/mq_dao_base.py +24 -13
- flowcept/commons/daos/mq_dao/mq_dao_kafka.py +18 -2
- flowcept/commons/flowcept_dataclasses/task_object.py +16 -21
- flowcept/commons/flowcept_dataclasses/workflow_object.py +9 -1
- flowcept/commons/task_data_preprocess.py +260 -60
- flowcept/commons/utils.py +25 -6
- flowcept/configs.py +41 -26
- flowcept/flowcept_api/flowcept_controller.py +73 -6
- flowcept/flowceptor/adapters/base_interceptor.py +11 -5
- flowcept/flowceptor/consumers/agent/base_agent_context_manager.py +25 -1
- flowcept/flowceptor/consumers/base_consumer.py +4 -0
- flowcept/flowceptor/consumers/consumer_utils.py +5 -4
- flowcept/flowceptor/consumers/document_inserter.py +2 -2
- flowcept/flowceptor/telemetry_capture.py +5 -2
- flowcept/instrumentation/flowcept_agent_task.py +294 -0
- flowcept/instrumentation/flowcept_decorator.py +43 -0
- flowcept/instrumentation/flowcept_loop.py +3 -3
- flowcept/instrumentation/flowcept_task.py +64 -24
- flowcept/instrumentation/flowcept_torch.py +5 -5
- flowcept/instrumentation/task_capture.py +83 -6
- flowcept/version.py +1 -1
- {flowcept-0.8.11.dist-info → flowcept-0.8.12.dist-info}/METADATA +42 -14
- {flowcept-0.8.11.dist-info → flowcept-0.8.12.dist-info}/RECORD +50 -36
- resources/sample_settings.yaml +12 -4
- flowcept/flowceptor/adapters/agents/__init__.py +0 -1
- flowcept/flowceptor/adapters/agents/agents_utils.py +0 -89
- flowcept/flowceptor/adapters/agents/flowcept_agent.py +0 -292
- flowcept/flowceptor/adapters/agents/flowcept_llm_prov_capture.py +0 -186
- flowcept/flowceptor/consumers/agent/flowcept_agent_context_manager.py +0 -145
- flowcept/flowceptor/consumers/agent/flowcept_qa_manager.py +0 -112
- {flowcept-0.8.11.dist-info → flowcept-0.8.12.dist-info}/WHEEL +0 -0
- {flowcept-0.8.11.dist-info → flowcept-0.8.12.dist-info}/entry_points.txt +0 -0
- {flowcept-0.8.11.dist-info → flowcept-0.8.12.dist-info}/licenses/LICENSE +0 -0
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
import streamlit as st
|
|
2
|
+
from flowcept.agents.gui import AI, PAGE_TITLE
|
|
3
|
+
from flowcept.agents.gui.gui_utils import (
|
|
4
|
+
query_agent,
|
|
5
|
+
display_ai_msg,
|
|
6
|
+
display_ai_msg_from_tool,
|
|
7
|
+
display_df_tool_response,
|
|
8
|
+
)
|
|
9
|
+
|
|
10
|
+
from flowcept.agents.tools.in_memory_queries.in_memory_queries_tools import (
|
|
11
|
+
generate_result_df,
|
|
12
|
+
generate_plot_code,
|
|
13
|
+
run_df_code,
|
|
14
|
+
)
|
|
15
|
+
|
|
16
|
+
st.set_page_config(page_title=PAGE_TITLE, page_icon=AI)
|
|
17
|
+
st.title(PAGE_TITLE)
|
|
18
|
+
|
|
19
|
+
GREETING = (
|
|
20
|
+
"Hi, there! I'm a **Workflow Provenance Specialist**.\n\n"
|
|
21
|
+
"I am tracking workflow executions and I can:\n"
|
|
22
|
+
"- 🔍 Analyze running workflows\n"
|
|
23
|
+
"- 📊 Plot graphs\n"
|
|
24
|
+
"- 🤖 Answer general questions about provenance data\n\n"
|
|
25
|
+
"How can I help you today?"
|
|
26
|
+
)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
display_ai_msg(GREETING)
|
|
30
|
+
|
|
31
|
+
# if "chat_history" not in st.session_state:
|
|
32
|
+
# st.session_state.chat_history = [{"role": "system", "content":GREETING}]
|
|
33
|
+
#
|
|
34
|
+
# for msg in st.session_state.chat_history:
|
|
35
|
+
# with st.chat_message(msg["role"], avatar=AI):
|
|
36
|
+
# st.markdown(msg["content"])
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def main():
|
|
40
|
+
"""Main Streamlit Function."""
|
|
41
|
+
user_input = st.chat_input("Send a message")
|
|
42
|
+
st.caption("💡 Tip: Ask about workflow metrics, generate plots, or summarize data.")
|
|
43
|
+
|
|
44
|
+
if user_input:
|
|
45
|
+
# st.session_state.chat_history.append({"role": "human", "content": user_input})
|
|
46
|
+
|
|
47
|
+
with st.chat_message("human"):
|
|
48
|
+
st.markdown(user_input)
|
|
49
|
+
|
|
50
|
+
try:
|
|
51
|
+
with st.spinner("🤖 Thinking..."):
|
|
52
|
+
tool_result = query_agent(user_input)
|
|
53
|
+
print(tool_result)
|
|
54
|
+
|
|
55
|
+
if tool_result.result_is_str():
|
|
56
|
+
display_ai_msg_from_tool(tool_result)
|
|
57
|
+
elif tool_result.is_success_dict():
|
|
58
|
+
tool_name = tool_result.tool_name
|
|
59
|
+
if tool_name in [generate_result_df.__name__, generate_plot_code.__name__, run_df_code.__name__]:
|
|
60
|
+
display_df_tool_response(tool_result)
|
|
61
|
+
else:
|
|
62
|
+
display_ai_msg(f"⚠️ Received unexpected response from agent: {tool_result}")
|
|
63
|
+
st.stop()
|
|
64
|
+
else:
|
|
65
|
+
display_df_tool_response(tool_result)
|
|
66
|
+
# display_ai_msg(f"⚠️ Received unexpected response from agent: {tool_result}")
|
|
67
|
+
st.stop()
|
|
68
|
+
|
|
69
|
+
except Exception as e:
|
|
70
|
+
display_ai_msg(f"❌ Error talking to MCP agent:\n\n```text\n{e}\n```")
|
|
71
|
+
st.stop()
|
|
72
|
+
|
|
73
|
+
# st.session_state.chat_history.append({"role": "system", "content": agent_reply})
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
main()
|
|
@@ -0,0 +1,239 @@
|
|
|
1
|
+
import io
|
|
2
|
+
import json
|
|
3
|
+
|
|
4
|
+
import streamlit as st
|
|
5
|
+
from flowcept.agents import prompt_handler
|
|
6
|
+
from flowcept.agents.agent_client import run_tool
|
|
7
|
+
from flowcept.agents.agents_utils import ToolResult
|
|
8
|
+
import pandas as pd
|
|
9
|
+
|
|
10
|
+
from flowcept.agents.gui import AI
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def query_agent(user_input: str) -> ToolResult:
|
|
14
|
+
"""
|
|
15
|
+
Send a user query to the agent and parse the response.
|
|
16
|
+
|
|
17
|
+
This function forwards the user input to the registered prompt handler
|
|
18
|
+
via ``run_tool``. The raw string response is then parsed into a
|
|
19
|
+
``ToolResult`` for structured handling of success and error cases.
|
|
20
|
+
|
|
21
|
+
Parameters
|
|
22
|
+
----------
|
|
23
|
+
user_input : str
|
|
24
|
+
The text query provided by the user.
|
|
25
|
+
|
|
26
|
+
Returns
|
|
27
|
+
-------
|
|
28
|
+
ToolResult
|
|
29
|
+
- ``code=400`` if the agent call fails.
|
|
30
|
+
- ``code=404`` if the agent response could not be parsed.
|
|
31
|
+
- ``code=499`` if JSON parsing fails.
|
|
32
|
+
- Otherwise, the parsed ``ToolResult`` object from the agent.
|
|
33
|
+
|
|
34
|
+
Examples
|
|
35
|
+
--------
|
|
36
|
+
>>> result = query_agent("Summarize the latest report.")
|
|
37
|
+
>>> if result.is_success():
|
|
38
|
+
... print(result.result)
|
|
39
|
+
"""
|
|
40
|
+
try:
|
|
41
|
+
response_str = run_tool(prompt_handler.__name__, kwargs={"message": user_input})[0]
|
|
42
|
+
except Exception as e:
|
|
43
|
+
return ToolResult(code=400, result=f"Failed to communicate with the Agent. Error: {e}")
|
|
44
|
+
try:
|
|
45
|
+
tool_result = ToolResult(**json.loads(response_str))
|
|
46
|
+
if tool_result is None:
|
|
47
|
+
ToolResult(code=404, result=f"Could not parse agent output:\n{response_str}")
|
|
48
|
+
return tool_result
|
|
49
|
+
except Exception as e:
|
|
50
|
+
return ToolResult(code=499, result=f"Failed to parse agent output:\n{response_str}.\n\nError: {e}")
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def display_ai_msg(msg: str):
|
|
54
|
+
"""
|
|
55
|
+
Display an AI message in the Streamlit chat interface.
|
|
56
|
+
|
|
57
|
+
This function creates a new chat message block with the "AI" role and
|
|
58
|
+
renders the given string as Markdown.
|
|
59
|
+
|
|
60
|
+
Parameters
|
|
61
|
+
----------
|
|
62
|
+
msg : str
|
|
63
|
+
The AI message to display.
|
|
64
|
+
|
|
65
|
+
Returns
|
|
66
|
+
-------
|
|
67
|
+
str
|
|
68
|
+
The same message string, useful for chaining or logging.
|
|
69
|
+
|
|
70
|
+
Examples
|
|
71
|
+
--------
|
|
72
|
+
>>> display_ai_msg("Hello! How can I help you today?")
|
|
73
|
+
"""
|
|
74
|
+
with st.chat_message("AI", avatar=AI):
|
|
75
|
+
st.markdown(msg)
|
|
76
|
+
return msg
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
def display_ai_msg_from_tool(tool_result: ToolResult):
|
|
80
|
+
"""
|
|
81
|
+
Display an AI message based on a ToolResult.
|
|
82
|
+
|
|
83
|
+
This function inspects the ``ToolResult`` to determine whether it
|
|
84
|
+
represents an error or a normal response. It then displays the
|
|
85
|
+
corresponding message in the Streamlit chat with the "AI" role.
|
|
86
|
+
|
|
87
|
+
Parameters
|
|
88
|
+
----------
|
|
89
|
+
tool_result : ToolResult
|
|
90
|
+
The tool result containing the agent's reply or error.
|
|
91
|
+
|
|
92
|
+
Returns
|
|
93
|
+
-------
|
|
94
|
+
str
|
|
95
|
+
The final message displayed in the chat.
|
|
96
|
+
|
|
97
|
+
Notes
|
|
98
|
+
-----
|
|
99
|
+
- If the result indicates an error (4xx codes), the message is shown in
|
|
100
|
+
a formatted error block with the error code.
|
|
101
|
+
- Otherwise, the raw result is displayed as Markdown.
|
|
102
|
+
|
|
103
|
+
Examples
|
|
104
|
+
--------
|
|
105
|
+
>>> res = ToolResult(code=301, result="Here is the summary you requested.")
|
|
106
|
+
>>> display_ai_msg_from_tool(res)
|
|
107
|
+
|
|
108
|
+
>>> err = ToolResult(code=405, result="Invalid JSON response")
|
|
109
|
+
>>> display_ai_msg_from_tool(err)
|
|
110
|
+
"""
|
|
111
|
+
has_error = tool_result.is_error_string()
|
|
112
|
+
with st.chat_message("AI", avatar=AI):
|
|
113
|
+
if has_error:
|
|
114
|
+
agent_reply = (
|
|
115
|
+
f"❌ Agent encountered an error, code {tool_result.code}:\n\n```text\n{tool_result.result}\n```"
|
|
116
|
+
)
|
|
117
|
+
else:
|
|
118
|
+
agent_reply = tool_result.result
|
|
119
|
+
|
|
120
|
+
st.markdown(agent_reply)
|
|
121
|
+
|
|
122
|
+
return agent_reply
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def display_df_tool_response(tool_result: ToolResult):
|
|
126
|
+
r"""
|
|
127
|
+
Display the DataFrame contained in a ToolResult.
|
|
128
|
+
|
|
129
|
+
This function extracts and displays the DataFrame (if present) from a
|
|
130
|
+
``ToolResult`` object, typically after executing a query or code
|
|
131
|
+
generation tool. It is intended for interactive use in environments
|
|
132
|
+
where DataFrame output should be visualized or printed.
|
|
133
|
+
|
|
134
|
+
Parameters
|
|
135
|
+
----------
|
|
136
|
+
tool_result : ToolResult
|
|
137
|
+
The tool result object containing the output of a previous operation.
|
|
138
|
+
Expected to include a CSV-formatted DataFrame string in its ``result``
|
|
139
|
+
field when ``code`` indicates success.
|
|
140
|
+
|
|
141
|
+
Notes
|
|
142
|
+
-----
|
|
143
|
+
- If the result does not contain a DataFrame, the function may print or
|
|
144
|
+
display an error message.
|
|
145
|
+
- The display method may vary depending on the environment (e.g., console,
|
|
146
|
+
Streamlit, or notebook).
|
|
147
|
+
|
|
148
|
+
Examples
|
|
149
|
+
--------
|
|
150
|
+
>>> result = ToolResult(code=301, result={"result_df": "col1,col2\\n1,2\\n3,4"})
|
|
151
|
+
>>> display_df_tool_response(result)
|
|
152
|
+
col1 col2
|
|
153
|
+
0 1 2
|
|
154
|
+
1 3 4
|
|
155
|
+
"""
|
|
156
|
+
result_dict = tool_result.result
|
|
157
|
+
result_code = result_dict.get("result_code", "")
|
|
158
|
+
result_df_str = result_dict.get("result_df", "").strip()
|
|
159
|
+
|
|
160
|
+
summary = result_dict.get("summary", "")
|
|
161
|
+
summary_error = result_dict.get("summary_error", "")
|
|
162
|
+
|
|
163
|
+
plot_code = result_dict.get("plot_code", "")
|
|
164
|
+
with st.chat_message("AI", avatar=AI):
|
|
165
|
+
st.markdown("📊 Here's the code:")
|
|
166
|
+
st.markdown(f"```python\n{result_code}")
|
|
167
|
+
print(result_code)
|
|
168
|
+
|
|
169
|
+
try:
|
|
170
|
+
df = pd.read_csv(io.StringIO(result_df_str))
|
|
171
|
+
print("The result is a df")
|
|
172
|
+
if not df.empty:
|
|
173
|
+
st.dataframe(df, hide_index=False)
|
|
174
|
+
print("Columns", str(df.columns))
|
|
175
|
+
print("Number of columns", len(df.columns))
|
|
176
|
+
else:
|
|
177
|
+
st.text("⚠️ Result DataFrame is empty.")
|
|
178
|
+
except Exception as e:
|
|
179
|
+
st.markdown(f"❌ {e}")
|
|
180
|
+
return
|
|
181
|
+
|
|
182
|
+
if plot_code:
|
|
183
|
+
st.markdown("Here's the plot code:")
|
|
184
|
+
st.markdown(f"```python\n{plot_code}")
|
|
185
|
+
st.markdown("📊 Here's the plot:")
|
|
186
|
+
try:
|
|
187
|
+
exec_st_plot_code(plot_code, df, st)
|
|
188
|
+
except Exception as e:
|
|
189
|
+
st.markdown(f"❌ {e}")
|
|
190
|
+
|
|
191
|
+
if summary:
|
|
192
|
+
st.markdown("📝 Summary:")
|
|
193
|
+
st.markdown(summary)
|
|
194
|
+
elif summary_error:
|
|
195
|
+
st.markdown(f"⚠️ Encountered this error when summarizing the result dataframe:\n```text\n{summary_error}")
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
def exec_st_plot_code(code, result_df, st_module):
|
|
199
|
+
"""
|
|
200
|
+
Execute plotting code dynamically with a given DataFrame and plotting modules.
|
|
201
|
+
|
|
202
|
+
This function runs a block of Python code (typically generated by an LLM)
|
|
203
|
+
to produce visualizations. It injects the provided DataFrame and plotting
|
|
204
|
+
libraries into the execution context, allowing the code to reference them
|
|
205
|
+
directly.
|
|
206
|
+
|
|
207
|
+
Parameters
|
|
208
|
+
----------
|
|
209
|
+
code : str
|
|
210
|
+
The Python code to execute, expected to contain plotting logic.
|
|
211
|
+
result_df : pandas.DataFrame
|
|
212
|
+
The DataFrame to be used within the plotting code (available as ``result``).
|
|
213
|
+
st_module : module
|
|
214
|
+
The Streamlit module (``st``) to be used within the plotting code.
|
|
215
|
+
|
|
216
|
+
Notes
|
|
217
|
+
-----
|
|
218
|
+
- The execution context includes:
|
|
219
|
+
- ``result`` : the provided DataFrame.
|
|
220
|
+
- ``st`` : the given Streamlit module.
|
|
221
|
+
- ``plt`` : ``matplotlib.pyplot`` for standard plotting.
|
|
222
|
+
- ``alt`` : ``altair`` for declarative plotting.
|
|
223
|
+
- The function uses Python's built-in ``exec``; malformed or unsafe code
|
|
224
|
+
may raise exceptions or cause side effects.
|
|
225
|
+
- Designed primarily for controlled scenarios such as running generated
|
|
226
|
+
plotting code inside an application.
|
|
227
|
+
|
|
228
|
+
Examples
|
|
229
|
+
--------
|
|
230
|
+
>>> import streamlit as st
|
|
231
|
+
>>> df = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
|
|
232
|
+
>>> code = "st.line_chart(result)"
|
|
233
|
+
>>> exec_st_plot_code(code, df, st)
|
|
234
|
+
"""
|
|
235
|
+
print("Plot code \n", code)
|
|
236
|
+
exec(
|
|
237
|
+
code,
|
|
238
|
+
{"result": result_df, "st": st_module, "plt": __import__("matplotlib.pyplot"), "alt": __import__("altair")},
|
|
239
|
+
)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""LLMs subpackage."""
|
|
@@ -0,0 +1,139 @@
|
|
|
1
|
+
import requests
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
class ClaudeOnGCPLLM:
|
|
5
|
+
"""
|
|
6
|
+
ClaudeOnGCPLLM is a wrapper for invoking Anthropic's Claude models
|
|
7
|
+
hosted on Google Cloud Vertex AI. It handles authentication, request
|
|
8
|
+
payload construction, and response parsing for text generation.
|
|
9
|
+
|
|
10
|
+
Parameters
|
|
11
|
+
----------
|
|
12
|
+
project_id : str
|
|
13
|
+
Google Cloud project ID used for Vertex AI requests.
|
|
14
|
+
google_token_auth : str
|
|
15
|
+
Bearer token for Google Cloud authentication.
|
|
16
|
+
location : str, default="us-east5"
|
|
17
|
+
Vertex AI location where the Claude model is hosted.
|
|
18
|
+
model_id : str, default="claude-opus-4"
|
|
19
|
+
Identifier of the Claude model to use.
|
|
20
|
+
anthropic_version : str, default="vertex-2023-10-16"
|
|
21
|
+
API version of Anthropic's Claude model on Vertex AI.
|
|
22
|
+
temperature : float, default=0.5
|
|
23
|
+
Sampling temperature controlling randomness of output.
|
|
24
|
+
max_tokens : int, default=512
|
|
25
|
+
Maximum number of tokens to generate in the response.
|
|
26
|
+
top_p : float, default=0.95
|
|
27
|
+
Nucleus sampling parameter; restricts tokens to a top cumulative probability.
|
|
28
|
+
top_k : int, default=1
|
|
29
|
+
Top-k sampling parameter; restricts tokens to the top-k most likely options.
|
|
30
|
+
|
|
31
|
+
Attributes
|
|
32
|
+
----------
|
|
33
|
+
url : str
|
|
34
|
+
Full REST endpoint URL for the Claude model on Vertex AI.
|
|
35
|
+
headers : dict
|
|
36
|
+
HTTP headers including authentication and content type.
|
|
37
|
+
temperature : float
|
|
38
|
+
Current temperature value used in requests.
|
|
39
|
+
max_tokens : int
|
|
40
|
+
Maximum number of tokens configured for output.
|
|
41
|
+
top_p : float
|
|
42
|
+
Probability cutoff for nucleus sampling.
|
|
43
|
+
top_k : int
|
|
44
|
+
Cutoff for top-k sampling.
|
|
45
|
+
|
|
46
|
+
Examples
|
|
47
|
+
--------
|
|
48
|
+
>>> llm = ClaudeOnGCPLLM(project_id="my-gcp-project", google_token_auth="ya29.a0...")
|
|
49
|
+
>>> response = llm.invoke("Write a poem about the sunrise.")
|
|
50
|
+
>>> print(response)
|
|
51
|
+
"A golden light spills across the horizon..."
|
|
52
|
+
"""
|
|
53
|
+
|
|
54
|
+
def __init__(
|
|
55
|
+
self,
|
|
56
|
+
project_id: str,
|
|
57
|
+
google_token_auth: str,
|
|
58
|
+
location: str = "us-east5",
|
|
59
|
+
model_id: str = "claude-opus-4",
|
|
60
|
+
anthropic_version: str = "vertex-2023-10-16",
|
|
61
|
+
temperature: float = 0.5,
|
|
62
|
+
max_tokens: int = 512,
|
|
63
|
+
top_p: float = 0.95,
|
|
64
|
+
top_k: int = 1,
|
|
65
|
+
):
|
|
66
|
+
self.project_id = project_id
|
|
67
|
+
self.location = location
|
|
68
|
+
self.model_id = model_id
|
|
69
|
+
self.anthropic_version = anthropic_version
|
|
70
|
+
self.endpoint = f"{location}-aiplatform.googleapis.com"
|
|
71
|
+
self.temperature = temperature
|
|
72
|
+
self.max_tokens = max_tokens
|
|
73
|
+
self.top_p = top_p
|
|
74
|
+
self.top_k = top_k
|
|
75
|
+
|
|
76
|
+
self.url = (
|
|
77
|
+
f"https://{self.endpoint}/v1/projects/{self.project_id}/locations/{self.location}"
|
|
78
|
+
f"/publishers/anthropic/models/{self.model_id}:rawPredict"
|
|
79
|
+
)
|
|
80
|
+
self.headers = {
|
|
81
|
+
"Authorization": f"Bearer {google_token_auth}",
|
|
82
|
+
"Content-Type": "application/json; charset=utf-8",
|
|
83
|
+
}
|
|
84
|
+
|
|
85
|
+
def invoke(self, prompt: str, **kwargs) -> str:
|
|
86
|
+
"""
|
|
87
|
+
Invoke the Claude model with a user prompt.
|
|
88
|
+
|
|
89
|
+
This method sends a prompt to the configured Claude model via Google
|
|
90
|
+
Cloud Vertex AI, waits for a response, and returns the generated text.
|
|
91
|
+
|
|
92
|
+
Parameters
|
|
93
|
+
----------
|
|
94
|
+
prompt : str
|
|
95
|
+
The user input to send to the Claude model.
|
|
96
|
+
**kwargs : dict, optional
|
|
97
|
+
Additional keyword arguments (currently unused, kept for extensibility).
|
|
98
|
+
|
|
99
|
+
Returns
|
|
100
|
+
-------
|
|
101
|
+
str
|
|
102
|
+
The generated text from the Claude model.
|
|
103
|
+
|
|
104
|
+
Raises
|
|
105
|
+
------
|
|
106
|
+
RuntimeError
|
|
107
|
+
If the Claude API call fails with a non-200 status code.
|
|
108
|
+
|
|
109
|
+
Examples
|
|
110
|
+
--------
|
|
111
|
+
>>> llm = ClaudeOnGCPLLM(project_id="my-gcp-project", google_token_auth="ya29.a0...")
|
|
112
|
+
>>> llm.invoke("Summarize the plot of Hamlet in two sentences.")
|
|
113
|
+
"Hamlet seeks to avenge his father’s death, feigns madness, and struggles with indecision.
|
|
114
|
+
Ultimately, nearly all the major characters perish, including Hamlet himself."
|
|
115
|
+
"""
|
|
116
|
+
payload = {
|
|
117
|
+
"anthropic_version": self.anthropic_version,
|
|
118
|
+
"stream": False,
|
|
119
|
+
"max_tokens": self.max_tokens,
|
|
120
|
+
"temperature": self.temperature,
|
|
121
|
+
"top_p": self.top_p,
|
|
122
|
+
"top_k": self.top_k,
|
|
123
|
+
"messages": [
|
|
124
|
+
{
|
|
125
|
+
"role": "user",
|
|
126
|
+
"content": [{"type": "text", "text": prompt}],
|
|
127
|
+
}
|
|
128
|
+
],
|
|
129
|
+
}
|
|
130
|
+
|
|
131
|
+
response = requests.post(self.url, headers=self.headers, json=payload)
|
|
132
|
+
|
|
133
|
+
if response.status_code != 200:
|
|
134
|
+
raise RuntimeError(f"Claude request failed: {response.status_code} {response.text}")
|
|
135
|
+
|
|
136
|
+
response_json = response.json()
|
|
137
|
+
|
|
138
|
+
# Return the text of the first content block
|
|
139
|
+
return response_json["content"][0]["text"]
|
|
@@ -0,0 +1,119 @@
|
|
|
1
|
+
from google import genai
|
|
2
|
+
from google.genai import types
|
|
3
|
+
import os
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class Gemini25LLM:
|
|
7
|
+
"""
|
|
8
|
+
Gemini25LLM is a lightweight wrapper around Google's Gemini 2.5 models
|
|
9
|
+
for text generation. It simplifies configuration and provides a unified
|
|
10
|
+
interface for invoking LLM completions with or without streaming.
|
|
11
|
+
|
|
12
|
+
Parameters
|
|
13
|
+
----------
|
|
14
|
+
project_id : str
|
|
15
|
+
Google Cloud project ID for authentication.
|
|
16
|
+
location : str, default="us-east5"
|
|
17
|
+
Vertex AI location where the model is hosted.
|
|
18
|
+
model : str, default="gemini-2.5-flash-lite"
|
|
19
|
+
The Gemini model to use (e.g., "gemini-2.5-flash", "gemini-2.5-pro").
|
|
20
|
+
temperature : float, default=0.7
|
|
21
|
+
Sampling temperature for controlling output randomness.
|
|
22
|
+
top_p : float, default=0.95
|
|
23
|
+
Nucleus sampling parameter; limits tokens to the top cumulative probability.
|
|
24
|
+
max_output_tokens : int, default=2048
|
|
25
|
+
Maximum number of tokens to generate in the response.
|
|
26
|
+
stream : bool, default=False
|
|
27
|
+
Whether to return responses incrementally (streaming) or as a single string.
|
|
28
|
+
|
|
29
|
+
Attributes
|
|
30
|
+
----------
|
|
31
|
+
model_name : str
|
|
32
|
+
Name of the Gemini model used for generation.
|
|
33
|
+
client : genai.Client
|
|
34
|
+
Underlying Google GenAI client instance.
|
|
35
|
+
config : types.GenerateContentConfig
|
|
36
|
+
Default generation configuration for the model.
|
|
37
|
+
stream : bool
|
|
38
|
+
Indicates whether streaming responses are enabled.
|
|
39
|
+
|
|
40
|
+
Examples
|
|
41
|
+
--------
|
|
42
|
+
Create a client and run a simple query:
|
|
43
|
+
|
|
44
|
+
>>> llm = Gemini25LLM(project_id="my-gcp-project")
|
|
45
|
+
>>> response = llm.invoke("Write a haiku about the ocean.")
|
|
46
|
+
>>> print(response)
|
|
47
|
+
"Blue waves rise and fall / endless dance beneath the sky / whispers of the deep"
|
|
48
|
+
"""
|
|
49
|
+
|
|
50
|
+
def __init__(
|
|
51
|
+
self,
|
|
52
|
+
project_id: str,
|
|
53
|
+
location: str = "us-east5",
|
|
54
|
+
model: str = "gemini-2.5-flash-lite",
|
|
55
|
+
temperature: float = 0.7,
|
|
56
|
+
top_p: float = 0.95,
|
|
57
|
+
max_output_tokens: int = 2048,
|
|
58
|
+
stream: bool = False,
|
|
59
|
+
):
|
|
60
|
+
self.model_name = model
|
|
61
|
+
os.environ["GOOGLE_CLOUD_PROJECT"] = project_id
|
|
62
|
+
self.stream = stream
|
|
63
|
+
self.client = genai.Client(vertexai=True, project=project_id, location=location)
|
|
64
|
+
self.config = types.GenerateContentConfig(
|
|
65
|
+
temperature=temperature,
|
|
66
|
+
top_p=top_p,
|
|
67
|
+
max_output_tokens=max_output_tokens,
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
def invoke(self, prompt: str, **kwargs) -> str:
|
|
71
|
+
r"""
|
|
72
|
+
Invoke the Gemini LLM with a user prompt.
|
|
73
|
+
|
|
74
|
+
This method sends the prompt to the configured Gemini model and returns
|
|
75
|
+
the generated text. It supports both streaming and non-streaming modes.
|
|
76
|
+
|
|
77
|
+
Parameters
|
|
78
|
+
----------
|
|
79
|
+
prompt : str
|
|
80
|
+
The input text prompt to send to the model.
|
|
81
|
+
**kwargs : dict, optional
|
|
82
|
+
Additional arguments (currently unused, kept for extensibility).
|
|
83
|
+
|
|
84
|
+
Returns
|
|
85
|
+
-------
|
|
86
|
+
str
|
|
87
|
+
The generated text response from the model. In streaming mode,
|
|
88
|
+
partial outputs are concatenated and returned as a single string.
|
|
89
|
+
|
|
90
|
+
Examples
|
|
91
|
+
--------
|
|
92
|
+
Basic invocation:
|
|
93
|
+
|
|
94
|
+
>>> llm = Gemini25LLM(project_id="my-gcp-project")
|
|
95
|
+
>>> llm.invoke("Explain quantum entanglement in simple terms.")
|
|
96
|
+
"A phenomenon where particles remain connected so that the state of one..."
|
|
97
|
+
|
|
98
|
+
Streaming invocation:
|
|
99
|
+
|
|
100
|
+
>>> llm = Gemini25LLM(project_id="my-gcp-project", stream=True)
|
|
101
|
+
>>> llm.invoke("List five creative startup ideas.")
|
|
102
|
+
"1. AI gardening assistant\n2. Virtual museum curator\n..."
|
|
103
|
+
"""
|
|
104
|
+
contents = [types.Content(role="user", parts=[types.Part.from_text(text=prompt)])]
|
|
105
|
+
|
|
106
|
+
if self.stream:
|
|
107
|
+
stream = self.client.models.generate_content_stream(
|
|
108
|
+
model=self.model_name,
|
|
109
|
+
contents=contents,
|
|
110
|
+
config=self.config,
|
|
111
|
+
)
|
|
112
|
+
return "".join(chunk.text for chunk in stream if chunk.text)
|
|
113
|
+
else:
|
|
114
|
+
result = self.client.models.generate_content(
|
|
115
|
+
model=self.model_name,
|
|
116
|
+
contents=contents,
|
|
117
|
+
config=self.config,
|
|
118
|
+
)
|
|
119
|
+
return result.text
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Prompts subpackage."""
|
|
@@ -1,3 +1,6 @@
|
|
|
1
|
+
# flake8: noqa: E501
|
|
2
|
+
# flake8: noqa: D103
|
|
3
|
+
|
|
1
4
|
from mcp.server.fastmcp.prompts import base
|
|
2
5
|
|
|
3
6
|
BASE_ROLE = (
|
|
@@ -14,6 +17,21 @@ DATA_SCHEMA_PROMPT = (
|
|
|
14
17
|
|
|
15
18
|
QUESTION_PROMPT = "I am particularly more interested in the following question: %QUESTION%."
|
|
16
19
|
|
|
20
|
+
SMALL_TALK_PROMPT = "Act as a Workflow Provenance Specialist. I would like to interact with you, but please be concise and brief. This is my message:\n"
|
|
21
|
+
|
|
22
|
+
ROUTING_PROMPT = (
|
|
23
|
+
"You are a routing assistant for a provenance AI agent. "
|
|
24
|
+
"Given the following user message, classify it into one of the following routes:\n"
|
|
25
|
+
"- small_talk: if it's casual conversation or some random word (e.g., 'hausdn', 'a', hello, how are you, what can you do, what's your name)\n"
|
|
26
|
+
"- plot: if user is requesting plots (e.g., plot, chart, visualize)\n"
|
|
27
|
+
"- historical_prov_query: if the user wants to query historical provenance data\n"
|
|
28
|
+
"- in_context_query: if the user appears to ask questions about tasks or data in running workflow (or a workflow that ran recently) or if the user mentions the in-memory 'df' or a dataframe.\n"
|
|
29
|
+
"- in_chat_query: if the user appears to be asking about something that has said recently in this chat.\n"
|
|
30
|
+
"- unknown: if you don't know.\n"
|
|
31
|
+
"Respond with only the route label."
|
|
32
|
+
"User message is below:\n "
|
|
33
|
+
)
|
|
34
|
+
|
|
17
35
|
|
|
18
36
|
def get_question_prompt(question: str):
|
|
19
37
|
"""Generates a user prompt with the given question filled in."""
|