forked from phoenix-oss/llama-stack-mirror
5de6ed946e
* Query generators for rag query * use agent.inference_api instead of passing host/port again * drop classes for functions --------- Co-authored-by: Hardik Shah <hjshah@fb.com>
76 lines
2.2 KiB
Python
76 lines
2.2 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 typing import List
|
|
|
|
from jinja2 import Template
|
|
from llama_models.llama3.api import * # noqa: F403
|
|
|
|
|
|
from llama_toolchain.agentic_system.api import (
|
|
DefaultMemoryQueryGeneratorConfig,
|
|
LLMMemoryQueryGeneratorConfig,
|
|
MemoryQueryGenerator,
|
|
MemoryQueryGeneratorConfig,
|
|
)
|
|
from termcolor import cprint # noqa: F401
|
|
from llama_toolchain.inference.api import * # noqa: F403
|
|
|
|
|
|
async def generate_rag_query(
|
|
config: MemoryQueryGeneratorConfig,
|
|
messages: List[Message],
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Generates a query that will be used for
|
|
retrieving relevant information from the memory bank.
|
|
"""
|
|
if config.type == MemoryQueryGenerator.default.value:
|
|
query = await default_rag_query_generator(config, messages, **kwargs)
|
|
elif config.type == MemoryQueryGenerator.llm.value:
|
|
query = await llm_rag_query_generator(config, messages, **kwargs)
|
|
else:
|
|
raise NotImplementedError(f"Unsupported memory query generator {config.type}")
|
|
# cprint(f"Generated query >>>: {query}", color="green")
|
|
return query
|
|
|
|
|
|
async def default_rag_query_generator(
|
|
config: DefaultMemoryQueryGeneratorConfig,
|
|
messages: List[Message],
|
|
**kwargs,
|
|
):
|
|
return config.sep.join(interleaved_text_media_as_str(m.content) for m in messages)
|
|
|
|
|
|
async def llm_rag_query_generator(
|
|
config: LLMMemoryQueryGeneratorConfig,
|
|
messages: List[Message],
|
|
**kwargs,
|
|
):
|
|
assert "inference_api" in kwargs, "LLMRAGQueryGenerator needs inference_api"
|
|
inference_api = kwargs["inference_api"]
|
|
|
|
m_dict = {"messages": [m.model_dump() for m in messages]}
|
|
|
|
template = Template(config.template)
|
|
content = template.render(m_dict)
|
|
|
|
model = config.model
|
|
message = UserMessage(content=content)
|
|
response = inference_api.chat_completion(
|
|
ChatCompletionRequest(
|
|
model=model,
|
|
messages=[message],
|
|
stream=False,
|
|
)
|
|
)
|
|
|
|
async for chunk in response:
|
|
query = chunk.completion_message.content
|
|
|
|
return query
|