apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 20:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

User/collection in structured query & agent

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-08 17:38:10 +01:00
parent 8a01790936
commit d03eb194d4
8 changed files with 66 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

30
tests/integration/test_agent_structured_query_integration.py
View file

@ -59,8 +59,10 @@ class TestAgentStructuredQueryIntegration:
# Create agent request
request = AgentRequest(
question="I need to find all customers from New York. Use the structured query tool to get this information.",
user="test_user",
collection="test_collection"
state="",
group=[],
history=[],
user="test_user"
)
msg = MagicMock()
@ -140,8 +142,10 @@ Args: {
request = AgentRequest(
question="Find data from a table that doesn't exist using structured query.",
user="test_user",
collection="test_collection"
state="",
group=[],
history=[],
user="test_user"
)
msg = MagicMock()
@ -209,8 +213,10 @@ Args: {
request = AgentRequest(
question="First find all customers from California, then tell me how many orders they have made.",
user="test_user",
collection="test_collection"
state="",
group=[],
history=[],
user="test_user"
)
msg = MagicMock()
@ -295,8 +301,10 @@ Args: {
request = AgentRequest(
question="Query the sales data for recent transactions.",
user="test_user",
collection="test_collection"
state="",
group=[],
history=[],
user="test_user"
)
msg = MagicMock()
@ -390,8 +398,10 @@ Args: {
request = AgentRequest(
question="Get customer information and format it nicely.",
user="test_user",
collection="test_collection"
state="",
group=[],
history=[],
user="test_user"
)
msg = MagicMock()

10
trustgraph-base/trustgraph/api/flow.py
View file

@ -132,11 +132,15 @@ class FlowInstance:
input
)["response"]
def agent(self, question):
def agent(self, question, user="trustgraph", state="", group=None, history=None):
# The input consists of a question
# The input consists of a question and optional context
input = {
"question": question
"question": question,
"user": user,
"state": state,
"group": group or [],
"history": history or []
}
return self.request(

6
trustgraph-base/trustgraph/base/structured_query_client.py
View file

@ -2,10 +2,12 @@ from . request_response_spec import RequestResponse, RequestResponseSpec
from .. schema import StructuredQueryRequest, StructuredQueryResponse
class StructuredQueryClient(RequestResponse):
async def structured_query(self, question, timeout=600):
async def structured_query(self, question, user="trustgraph", collection="default", timeout=600):
resp = await self.request(
StructuredQueryRequest(
question = question
question = question,
user = user,
collection = collection
),
timeout=timeout
)

10
trustgraph-base/trustgraph/messaging/translators/agent.py
View file

@ -9,17 +9,19 @@ class AgentRequestTranslator(MessageTranslator):
def to_pulsar(self, data: Dict[str, Any]) -> AgentRequest:
return AgentRequest(
question=data["question"],
plan=data.get("plan", ""),
state=data.get("state", ""),
history=data.get("history", [])
group=data.get("group", []),
history=data.get("history", []),
user=data.get("user", "trustgraph")
)
def from_pulsar(self, obj: AgentRequest) -> Dict[str, Any]:
return {
"question": obj.question,
"plan": obj.plan,
"state": obj.state,
"history": obj.history
"group": obj.group,
"history": obj.history,
"user": obj.user
}

2
trustgraph-base/trustgraph/schema/services/agent.py
View file

@ -13,12 +13,14 @@ class AgentStep(Record):
action = String()
arguments = Map(String())
observation = String()
user = String() # User context for the step
class AgentRequest(Record):
question = String()
state = String()
group = Array(String())
history = Array(AgentStep())
user = String() # User context for multi-tenancy
class AgentResponse(Record):
answer = String()

4
trustgraph-cli/trustgraph/cli/invoke_agent.py
View file

@ -61,6 +61,10 @@ async def question(
"flow": flow_id,
"request": {
"question": question,
"user": user,
"state": state or "",
"group": [],
"history": []
}
})

16
trustgraph-flow/trustgraph/agent/react/service.py
View file

@ -253,12 +253,26 @@ class Processor(AgentService):
logger.debug("Call React")
# Create user-aware context wrapper that preserves the flow interface
# but adds user information for tools that need it
class UserAwareContext:
def __init__(self, flow, user):
self._flow = flow
self._user = user
def __call__(self, service_name):
client = self._flow(service_name)
# For structured query clients, store user context
if service_name == "structured-query-request":
client._current_user = self._user
return client
act = await temp_agent.react(
question = request.question,
history = history,
think = think,
observe = observe,
context = flow,
context = UserAwareContext(flow, request.user),
)
logger.debug(f"Action: {act}")

10
trustgraph-flow/trustgraph/agent/react/tools.py
View file

@ -87,9 +87,10 @@ class McpToolImpl:
# This tool implementation knows how to query structured data using natural language
class StructuredQueryImpl:
def __init__(self, context, collection=None):
def __init__(self, context, collection=None, user=None):
self.context = context
self.collection = collection # For multi-tenant scenarios
self.user = user # User context for multi-tenancy
@staticmethod
def get_arguments():
@ -105,8 +106,13 @@ class StructuredQueryImpl:
client = self.context("structured-query-request")
logger.debug("Structured query question...")
# Get user from client context if available, otherwise use instance user or default
user = getattr(client, '_current_user', self.user or "trustgraph")
result = await client.structured_query(
arguments.get("question")
question=arguments.get("question"),
user=user,
collection=self.collection or "default"
)
# Format the result for the agent