phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-08-03 09:21:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add memory bank attributes

Browse source
This commit is contained in:
Dinesh Yeduguru 2024-11-25 11:54:40 -08:00
parent b1a63df8cd
commit 6411007024
2 changed files with 12 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llama_stack/providers/inline/agents/meta_reference/agent_instance.py
View file

@ -356,10 +356,12 @@ class ChatAgent(ShieldRunnerMixin):
# TODO: find older context from the session and either replace it # TODO: find older context from the session and either replace it
# or append with a sliding window. this is really a very simplistic implementation # or append with a sliding window. this is really a very simplistic implementation
with tracing.span("retrieve_rag_context"): with tracing.span("retrieve_rag_context") as span:
rag_context, bank_ids = await self._retrieve_context( rag_context, bank_ids = await self._retrieve_context(
session_id, input_messages, attachments session_id, input_messages, attachments
) )
span.set_attribute("rag_context", rag_context)
span.set_attribute("bank_ids", bank_ids)
step_id = str(uuid.uuid4()) step_id = str(uuid.uuid4())
yield AgentTurnResponseStreamChunk( yield AgentTurnResponseStreamChunk(

11
llama_stack/providers/utils/telemetry/tracing.py
View file

@ -69,7 +69,7 @@ class TraceContext:
self.logger = logger self.logger = logger
self.trace_id = trace_id self.trace_id = trace_id
def push_span(self, name: str, attributes: Dict[str, Any] = None): def push_span(self, name: str, attributes: Dict[str, Any] = None) -> Span:
current_span = self.get_current_span() current_span = self.get_current_span()
span = Span( span = Span(
span_id=generate_short_uuid(), span_id=generate_short_uuid(),
@ -94,6 +94,7 @@ class TraceContext:
) )
self.spans.append(span) self.spans.append(span)
return span
def pop_span(self, status: SpanStatus = SpanStatus.OK): def pop_span(self, status: SpanStatus = SpanStatus.OK):
span = self.spans.pop() span = self.spans.pop()
@ -208,7 +209,7 @@ class SpanContextManager:
global CURRENT_TRACE_CONTEXT global CURRENT_TRACE_CONTEXT
context = CURRENT_TRACE_CONTEXT context = CURRENT_TRACE_CONTEXT
if context: if context:
context.push_span(self.name, self.attributes) self.span = context.push_span(self.name, self.attributes)
return self return self
def __exit__(self, exc_type, exc_value, traceback): def __exit__(self, exc_type, exc_value, traceback):
@ -217,6 +218,12 @@ class SpanContextManager:
if context: if context:
context.pop_span() context.pop_span()
def set_attribute(self, key: str, value: Any):
if self.span:
if self.span.attributes is None:
self.span.attributes = {}
self.span.attributes[key] = value
async def __aenter__(self): async def __aenter__(self):
return self.__enter__() return self.__enter__()