phoenix-oss/llama-stack
1
0
Fork 0
forked from phoenix-oss/llama-stack-mirror
Code Pull requests Releases Packages 2 Activity Actions

Respect passed in embedding model

Browse source
This commit is contained in:
Ashwin Bharambe 2024-09-24 14:40:28 -07:00
parent bda974e660
commit 00352bd251
2 changed files with 16 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

6
llama_stack/apis/memory/client.py
View file

@ -13,9 +13,9 @@ from typing import Any, Dict, List, Optional
import fire
import httpx
from termcolor import cprint
from llama_stack.distribution.datatypes import RemoteProviderConfig
from termcolor import cprint
from llama_stack.apis.memory import * # noqa: F403
from llama_stack.providers.utils.memory.file_utils import data_url_from_file
@ -120,7 +120,7 @@ async def run_main(host: str, port: int, stream: bool):
name="test_bank",
config=VectorMemoryBankConfig(
bank_id="test_bank",
embedding_model="dragon-roberta-query-2",
embedding_model="all-MiniLM-L6-v2",
chunk_size_in_tokens=512,
overlap_size_in_tokens=64,
),
@ -129,7 +129,7 @@ async def run_main(host: str, port: int, stream: bool):
retrieved_bank = await client.get_memory_bank(bank.bank_id)
assert retrieved_bank is not None
assert retrieved_bank.config.embedding_model == "dragon-roberta-query-2"
assert retrieved_bank.config.embedding_model == "all-MiniLM-L6-v2"
urls = [
"memory_optimizations.rst",