forked from phoenix-oss/llama-stack-mirror
fcd6449519
# What does this PR do? Change the Telemetry API to be able to support different use cases like returning traces for the UI and ability to export for Evals. Other changes: * Add a new trace_protocol decorator to decorate all our API methods so that any call to them will automatically get traced across all impls. * There is some issue with the decorator pattern of span creation when using async generators, where there are multiple yields with in the same context. I think its much more explicit by using the explicit context manager pattern using with. I moved the span creations in agent instance to be using with * Inject session id at the turn level, which should quickly give us all traces across turns for a given session Addresses #509 ## Test Plan ``` llama stack run /Users/dineshyv/.llama/distributions/llamastack-together/together-run.yaml PYTHONPATH=. python -m examples.agents.rag_with_memory_bank localhost 5000 curl -X POST 'http://localhost:5000/alpha/telemetry/query-traces' \ -H 'Content-Type: application/json' \ -d '{ "attribute_filters": [ { "key": "session_id", "op": "eq", "value": "dd667b87-ca4b-4d30-9265-5a0de318fc65" }], "limit": 100, "offset": 0, "order_by": ["start_time"] }' | jq . [ { "trace_id": "6902f54b83b4b48be18a6f422b13e16f", "root_span_id": "5f37b85543afc15a", "start_time": "2024-12-04T08:08:30.501587", "end_time": "2024-12-04T08:08:36.026463" }, { "trace_id": "92227dac84c0615ed741be393813fb5f", "root_span_id": "af7c5bb46665c2c8", "start_time": "2024-12-04T08:08:36.031170", "end_time": "2024-12-04T08:08:41.693301" }, { "trace_id": "7d578a6edac62f204ab479fba82f77b6", "root_span_id": "1d935e3362676896", "start_time": "2024-12-04T08:08:41.695204", "end_time": "2024-12-04T08:08:47.228016" }, { "trace_id": "dbd767d76991bc816f9f078907dc9ff2", "root_span_id": "f5a7ee76683b9602", "start_time": "2024-12-04T08:08:47.234578", "end_time": "2024-12-04T08:08:53.189412" } ] curl -X POST 'http://localhost:5000/alpha/telemetry/get-span-tree' \ -H 'Content-Type: application/json' \ -d '{ "span_id" : "6cceb4b48a156913", "max_depth": 2, "attributes_to_return": ["input"] }' | jq . % Total % Received % Xferd Average Speed Time Time Time Current Dload Upload Total Spent Left Speed 100 875 100 790 100 85 18462 1986 --:--:-- --:--:-- --:--:-- 20833 { "span_id": "6cceb4b48a156913", "trace_id": "dafa796f6aaf925f511c04cd7c67fdda", "parent_span_id": "892a66d726c7f990", "name": "retrieve_rag_context", "start_time": "2024-12-04T09:28:21.781995", "end_time": "2024-12-04T09:28:21.913352", "attributes": { "input": [ "{\"role\":\"system\",\"content\":\"You are a helpful assistant\"}", "{\"role\":\"user\",\"content\":\"What are the top 5 topics that were explained in the documentation? Only list succinct bullet points.\",\"context\":null}" ] }, "children": [ { "span_id": "1a2df181854064a8", "trace_id": "dafa796f6aaf925f511c04cd7c67fdda", "parent_span_id": "6cceb4b48a156913", "name": "MemoryRouter.query_documents", "start_time": "2024-12-04T09:28:21.787620", "end_time": "2024-12-04T09:28:21.906512", "attributes": { "input": null }, "children": [], "status": "ok" } ], "status": "ok" } ``` <img width="1677" alt="Screenshot 2024-12-04 at 9 42 56 AM" src="https://github.com/user-attachments/assets/4d3cea93-05ce-415a-93d9-4b1628631bf8">
151 lines
4 KiB
Python
151 lines
4 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
from enum import Enum
|
|
from typing import (
|
|
Annotated,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Protocol,
|
|
runtime_checkable,
|
|
Union,
|
|
)
|
|
|
|
from llama_models.schema_utils import json_schema_type, webmethod
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from llama_stack.apis.resource import Resource, ResourceType
|
|
from llama_stack.distribution.tracing import trace_protocol
|
|
|
|
|
|
@json_schema_type
|
|
class MemoryBankType(Enum):
|
|
vector = "vector"
|
|
keyvalue = "keyvalue"
|
|
keyword = "keyword"
|
|
graph = "graph"
|
|
|
|
|
|
# define params for each type of memory bank, this leads to a tagged union
|
|
# accepted as input from the API or from the config.
|
|
@json_schema_type
|
|
class VectorMemoryBankParams(BaseModel):
|
|
memory_bank_type: Literal[MemoryBankType.vector.value] = MemoryBankType.vector.value
|
|
embedding_model: str
|
|
chunk_size_in_tokens: int
|
|
overlap_size_in_tokens: Optional[int] = None
|
|
|
|
|
|
@json_schema_type
|
|
class KeyValueMemoryBankParams(BaseModel):
|
|
memory_bank_type: Literal[MemoryBankType.keyvalue.value] = (
|
|
MemoryBankType.keyvalue.value
|
|
)
|
|
|
|
|
|
@json_schema_type
|
|
class KeywordMemoryBankParams(BaseModel):
|
|
memory_bank_type: Literal[MemoryBankType.keyword.value] = (
|
|
MemoryBankType.keyword.value
|
|
)
|
|
|
|
|
|
@json_schema_type
|
|
class GraphMemoryBankParams(BaseModel):
|
|
memory_bank_type: Literal[MemoryBankType.graph.value] = MemoryBankType.graph.value
|
|
|
|
|
|
BankParams = Annotated[
|
|
Union[
|
|
VectorMemoryBankParams,
|
|
KeyValueMemoryBankParams,
|
|
KeywordMemoryBankParams,
|
|
GraphMemoryBankParams,
|
|
],
|
|
Field(discriminator="memory_bank_type"),
|
|
]
|
|
|
|
|
|
# Some common functionality for memory banks.
|
|
class MemoryBankResourceMixin(Resource):
|
|
type: Literal[ResourceType.memory_bank.value] = ResourceType.memory_bank.value
|
|
|
|
@property
|
|
def memory_bank_id(self) -> str:
|
|
return self.identifier
|
|
|
|
@property
|
|
def provider_memory_bank_id(self) -> str:
|
|
return self.provider_resource_id
|
|
|
|
|
|
@json_schema_type
|
|
class VectorMemoryBank(MemoryBankResourceMixin):
|
|
memory_bank_type: Literal[MemoryBankType.vector.value] = MemoryBankType.vector.value
|
|
embedding_model: str
|
|
chunk_size_in_tokens: int
|
|
overlap_size_in_tokens: Optional[int] = None
|
|
|
|
|
|
@json_schema_type
|
|
class KeyValueMemoryBank(MemoryBankResourceMixin):
|
|
memory_bank_type: Literal[MemoryBankType.keyvalue.value] = (
|
|
MemoryBankType.keyvalue.value
|
|
)
|
|
|
|
|
|
# TODO: KeyValue and Keyword are so similar in name, oof. Get a better naming convention.
|
|
@json_schema_type
|
|
class KeywordMemoryBank(MemoryBankResourceMixin):
|
|
memory_bank_type: Literal[MemoryBankType.keyword.value] = (
|
|
MemoryBankType.keyword.value
|
|
)
|
|
|
|
|
|
@json_schema_type
|
|
class GraphMemoryBank(MemoryBankResourceMixin):
|
|
memory_bank_type: Literal[MemoryBankType.graph.value] = MemoryBankType.graph.value
|
|
|
|
|
|
MemoryBank = Annotated[
|
|
Union[
|
|
VectorMemoryBank,
|
|
KeyValueMemoryBank,
|
|
KeywordMemoryBank,
|
|
GraphMemoryBank,
|
|
],
|
|
Field(discriminator="memory_bank_type"),
|
|
]
|
|
|
|
|
|
class MemoryBankInput(BaseModel):
|
|
memory_bank_id: str
|
|
params: BankParams
|
|
provider_memory_bank_id: Optional[str] = None
|
|
|
|
|
|
@runtime_checkable
|
|
@trace_protocol
|
|
class MemoryBanks(Protocol):
|
|
@webmethod(route="/memory-banks/list", method="GET")
|
|
async def list_memory_banks(self) -> List[MemoryBank]: ...
|
|
|
|
@webmethod(route="/memory-banks/get", method="GET")
|
|
async def get_memory_bank(self, memory_bank_id: str) -> Optional[MemoryBank]: ...
|
|
|
|
@webmethod(route="/memory-banks/register", method="POST")
|
|
async def register_memory_bank(
|
|
self,
|
|
memory_bank_id: str,
|
|
params: BankParams,
|
|
provider_id: Optional[str] = None,
|
|
provider_memory_bank_id: Optional[str] = None,
|
|
) -> MemoryBank: ...
|
|
|
|
@webmethod(route="/memory-banks/unregister", method="POST")
|
|
async def unregister_memory_bank(self, memory_bank_id: str) -> None: ...
|