Add patch for SSE event endpoint responses (#50)

This commit is contained in:
Dalton Flanagan 2024-09-03 23:40:31 -04:00 committed by GitHub
parent 0af81776c7
commit 35093c0b6f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 15 additions and 0 deletions

View file

@ -36,6 +36,11 @@ schema_utils.json_schema_type = json_schema_type
from llama_toolchain.stack import LlamaStack
STREAMING_ENDPOINTS = [
"/agentic_system/turn/create"
]
def patched_get_endpoint_functions(
endpoint: type, prefixes: List[str]
) -> Iterator[Tuple[str, str, str, Callable]]:
@ -75,6 +80,13 @@ def patched_get_endpoint_functions(
operations._get_endpoint_functions = patched_get_endpoint_functions
def patch_sse_stream_responses(spec: Specification):
for path, path_item in spec.document.paths.items():
if path in STREAMING_ENDPOINTS:
content = path_item.post.responses['200'].content.pop('application/json')
path_item.post.responses['200'].content['text/event-stream'] = content
def main(output_dir: str):
output_dir = Path(output_dir)
if not output_dir.exists():
@ -100,6 +112,9 @@ def main(output_dir: str):
),
),
)
patch_sse_stream_responses(spec)
with open(output_dir / "llama-stack-spec.yaml", "w", encoding="utf-8") as fp:
yaml.dump(spec.get_json(), fp, allow_unicode=True)

0
rfcs/openapi_generator/run_openapi_generator.sh Normal file → Executable file
View file