llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/impls/meta_reference/agents/tests/test_chat_agent.py

@@ -1,305 +0,0 @@
-# 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 AsyncIterator, List, Optional, Union
-
-import pytest
-
-from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.apis.safety import *  # noqa: F403
-from llama_stack.apis.agents import *  # noqa: F403
-
-from ..agents import (
-    AGENT_INSTANCES_BY_ID,
-    MetaReferenceAgentsImpl,
-    MetaReferenceInferenceConfig,
-)
-
-
-class MockInferenceAPI:
-    async def chat_completion(
-        self,
-        model: str,
-        messages: List[Message],
-        sampling_params: Optional[SamplingParams] = SamplingParams(),
-        tools: Optional[List[ToolDefinition]] = None,
-        tool_choice: Optional[ToolChoice] = None,
-        tool_prompt_format: Optional[ToolPromptFormat] = None,
-        stream: Optional[bool] = False,
-        logprobs: Optional[LogProbConfig] = None,
-    ) -> AsyncIterator[
-        Union[ChatCompletionResponseStreamChunk, ChatCompletionResponse]
-    ]:
-        if stream:
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type="start",
-                    delta="",
-                )
-            )
-
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type="progress",
-                    delta="AI is a fascinating field...",
-                )
-            )
-            # yield ChatCompletionResponseStreamChunk(
-            #     event=ChatCompletionResponseEvent(
-            #         event_type="progress",
-            #         delta=ToolCallDelta(
-            #             content=ToolCall(
-            #                 call_id="123",
-            #                 tool_name=BuiltinTool.brave_search.value,
-            #                 arguments={"query": "AI history"},
-            #             ),
-            #             parse_status="success",
-            #         ),
-            #     )
-            # )
-            yield ChatCompletionResponseStreamChunk(
-                event=ChatCompletionResponseEvent(
-                    event_type="complete",
-                    delta="",
-                    stop_reason="end_of_turn",
-                )
-            )
-        else:
-            yield ChatCompletionResponse(
-                completion_message=CompletionMessage(
-                    role="assistant", content="Mock response", stop_reason="end_of_turn"
-                ),
-                logprobs=[0.1, 0.2, 0.3] if logprobs else None,
-            )
-
-
-class MockSafetyAPI:
-    async def run_shield(
-        self, shield_type: str, messages: List[Message]
-    ) -> RunShieldResponse:
-        return RunShieldResponse(violation=None)
-
-
-class MockMemoryAPI:
-    def __init__(self):
-        self.memory_banks = {}
-        self.documents = {}
-
-    async def create_memory_bank(self, name, config, url=None):
-        bank_id = f"bank_{len(self.memory_banks)}"
-        bank = MemoryBank(bank_id, name, config, url)
-        self.memory_banks[bank_id] = bank
-        self.documents[bank_id] = {}
-        return bank
-
-    async def list_memory_banks(self):
-        return list(self.memory_banks.values())
-
-    async def get_memory_bank(self, bank_id):
-        return self.memory_banks.get(bank_id)
-
-    async def drop_memory_bank(self, bank_id):
-        if bank_id in self.memory_banks:
-            del self.memory_banks[bank_id]
-            del self.documents[bank_id]
-        return bank_id
-
-    async def insert_documents(self, bank_id, documents, ttl_seconds=None):
-        if bank_id not in self.documents:
-            raise ValueError(f"Bank {bank_id} not found")
-        for doc in documents:
-            self.documents[bank_id][doc.document_id] = doc
-
-    async def update_documents(self, bank_id, documents):
-        if bank_id not in self.documents:
-            raise ValueError(f"Bank {bank_id} not found")
-        for doc in documents:
-            if doc.document_id in self.documents[bank_id]:
-                self.documents[bank_id][doc.document_id] = doc
-
-    async def query_documents(self, bank_id, query, params=None):
-        if bank_id not in self.documents:
-            raise ValueError(f"Bank {bank_id} not found")
-        # Simple mock implementation: return all documents
-        chunks = [
-            {"content": doc.content, "token_count": 10, "document_id": doc.document_id}
-            for doc in self.documents[bank_id].values()
-        ]
-        scores = [1.0] * len(chunks)
-        return {"chunks": chunks, "scores": scores}
-
-    async def get_documents(self, bank_id, document_ids):
-        if bank_id not in self.documents:
-            raise ValueError(f"Bank {bank_id} not found")
-        return [
-            self.documents[bank_id][doc_id]
-            for doc_id in document_ids
-            if doc_id in self.documents[bank_id]
-        ]
-
-    async def delete_documents(self, bank_id, document_ids):
-        if bank_id not in self.documents:
-            raise ValueError(f"Bank {bank_id} not found")
-        for doc_id in document_ids:
-            self.documents[bank_id].pop(doc_id, None)
-
-
-@pytest.fixture
-def mock_inference_api():
-    return MockInferenceAPI()
-
-
-@pytest.fixture
-def mock_safety_api():
-    return MockSafetyAPI()
-
-
-@pytest.fixture
-def mock_memory_api():
-    return MockMemoryAPI()
-
-
-@pytest.fixture
-async def chat_agent(mock_inference_api, mock_safety_api, mock_memory_api):
-    impl = MetaReferenceAgentsImpl(
-        config=MetaReferenceInferenceConfig(),
-        inference_api=mock_inference_api,
-        safety_api=mock_safety_api,
-        memory_api=mock_memory_api,
-    )
-    await impl.initialize()
-
-    agent_config = AgentConfig(
-        model="test_model",
-        instructions="You are a helpful assistant.",
-        sampling_params=SamplingParams(),
-        tools=[
-            # SearchToolDefinition(
-            #     name="brave_search",
-            #     api_key="test_key",
-            # ),
-        ],
-        tool_choice=ToolChoice.auto,
-        enable_session_persistence=False,
-        input_shields=[],
-        output_shields=[],
-    )
-    response = await impl.create_agent(agent_config)
-    agent = AGENT_INSTANCES_BY_ID[response.agent_id]
-    return agent
-
-
-@pytest.mark.asyncio
-async def test_chat_agent_create_session(chat_agent):
-    session = chat_agent.create_session("Test Session")
-    assert session.session_name == "Test Session"
-    assert session.turns == []
-    assert session.session_id in chat_agent.sessions
-
-
-@pytest.mark.asyncio
-async def test_chat_agent_create_and_execute_turn(chat_agent):
-    session = chat_agent.create_session("Test Session")
-    request = AgentTurnCreateRequest(
-        agent_id="random",
-        session_id=session.session_id,
-        messages=[UserMessage(content="Hello")],
-    )
-
-    responses = []
-    async for response in chat_agent.create_and_execute_turn(request):
-        responses.append(response)
-
-    print(responses)
-    assert len(responses) > 0
-    assert len(responses) == 4  # TurnStart, StepStart, StepComplete, TurnComplete
-    assert responses[0].event.payload.turn_id is not None
-
-
-@pytest.mark.asyncio
-async def test_run_multiple_shields_wrapper(chat_agent):
-    messages = [UserMessage(content="Test message")]
-    shields = ["test_shield"]
-
-    responses = [
-        chunk
-        async for chunk in chat_agent.run_multiple_shields_wrapper(
-            turn_id="test_turn_id",
-            messages=messages,
-            shields=shields,
-            touchpoint="user-input",
-        )
-    ]
-
-    assert len(responses) == 2  # StepStart, StepComplete
-    assert responses[0].event.payload.step_type.value == "shield_call"
-    assert not responses[1].event.payload.step_details.response.is_violation
-
-
-@pytest.mark.asyncio
-@pytest.mark.skip(reason="Not yet implemented; need to mock out tool execution easily")
-async def test_chat_agent_complex_turn(chat_agent):
-    # Setup
-    session = chat_agent.create_session("Test Session")
-    request = AgentTurnCreateRequest(
-        agent_id="random",
-        session_id=session.session_id,
-        messages=[UserMessage(content="Tell me about AI and then use a tool.")],
-        stream=True,
-    )
-
-    # Execute the turn
-    responses = []
-    async for response in chat_agent.create_and_execute_turn(request):
-        responses.append(response)
-
-    # Assertions
-    assert len(responses) > 0
-
-    # Check for the presence of different step types
-    step_types = [
-        response.event.payload.step_type
-        for response in responses
-        if hasattr(response.event.payload, "step_type")
-    ]
-
-    assert "shield_call" in step_types, "Shield call step is missing"
-    assert "inference" in step_types, "Inference step is missing"
-    assert "tool_execution" in step_types, "Tool execution step is missing"
-
-    # Check for the presence of start and complete events
-    event_types = [
-        response.event.payload.event_type
-        for response in responses
-        if hasattr(response.event.payload, "event_type")
-    ]
-    assert "start" in event_types, "Start event is missing"
-    assert "complete" in event_types, "Complete event is missing"
-
-    # Check for the presence of tool call
-    tool_calls = [
-        response.event.payload.tool_call
-        for response in responses
-        if hasattr(response.event.payload, "tool_call")
-    ]
-    assert any(
-        tool_call
-        for tool_call in tool_calls
-        if tool_call and tool_call.content.get("name") == "memory"
-    ), "Memory tool call is missing"
-
-    # Check for the final turn complete event
-    assert any(
-        isinstance(response.event.payload, AgentTurnResponseTurnCompletePayload)
-        for response in responses
-    ), "Turn complete event is missing"
-
-    # Verify the turn was added to the session
-    assert len(session.turns) == 1, "Turn was not added to the session"
-    assert (
-        session.turns[0].input_messages == request.messages
-    ), "Input messages do not match"

llama_stack/providers/impls/meta_reference/agents/tools/base.py

@@ -1,20 +0,0 @@
-# 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 abc import ABC, abstractmethod
-from typing import List
-
-from llama_stack.apis.inference import Message
-
-
-class BaseTool(ABC):
-    @abstractmethod
-    def get_name(self) -> str:
-        raise NotImplementedError
-
-    @abstractmethod
-    async def run(self, messages: List[Message]) -> List[Message]:
-        raise NotImplementedError

llama_stack/providers/impls/meta_reference/agents/tools/builtin.py

@@ -1,375 +0,0 @@
-# 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.
-
-import json
-import re
-import tempfile
-
-from abc import abstractmethod
-from typing import List, Optional
-
-import requests
-from termcolor import cprint
-
-from .ipython_tool.code_execution import (
-    CodeExecutionContext,
-    CodeExecutionRequest,
-    CodeExecutor,
-    TOOLS_ATTACHMENT_KEY_REGEX,
-)
-
-from llama_stack.apis.inference import *  # noqa: F403
-from llama_stack.apis.agents import *  # noqa: F403
-
-from .base import BaseTool
-
-
-def interpret_content_as_attachment(content: str) -> Optional[Attachment]:
-    match = re.search(TOOLS_ATTACHMENT_KEY_REGEX, content)
-    if match:
-        snippet = match.group(1)
-        data = json.loads(snippet)
-        return Attachment(
-            content=URL(uri="file://" + data["filepath"]), mime_type=data["mimetype"]
-        )
-
-    return None
-
-
-class SingleMessageBuiltinTool(BaseTool):
-    async def run(self, messages: List[CompletionMessage]) -> List[ToolResponseMessage]:
-        assert len(messages) == 1, f"Expected single message, got {len(messages)}"
-
-        message = messages[0]
-        assert len(message.tool_calls) == 1, "Expected a single tool call"
-
-        tool_call = messages[0].tool_calls[0]
-
-        query = tool_call.arguments["query"]
-        response: str = await self.run_impl(query)
-
-        message = ToolResponseMessage(
-            call_id=tool_call.call_id,
-            tool_name=tool_call.tool_name,
-            content=response,
-        )
-        return [message]
-
-    @abstractmethod
-    async def run_impl(self, query: str) -> str:
-        raise NotImplementedError()
-
-
-class PhotogenTool(SingleMessageBuiltinTool):
-    def __init__(self, dump_dir: str) -> None:
-        self.dump_dir = dump_dir
-
-    def get_name(self) -> str:
-        return BuiltinTool.photogen.value
-
-    async def run_impl(self, query: str) -> str:
-        """
-        Implement this to give the model an ability to generate images.
-
-        Return:
-            info = {
-                "filepath": str(image_filepath),
-                "mimetype": "image/png",
-            }
-        """
-        raise NotImplementedError()
-
-
-class SearchTool(SingleMessageBuiltinTool):
-    def __init__(self, engine: SearchEngineType, api_key: str, **kwargs) -> None:
-        self.api_key = api_key
-        if engine == SearchEngineType.bing:
-            self.engine = BingSearch(api_key, **kwargs)
-        elif engine == SearchEngineType.brave:
-            self.engine = BraveSearch(api_key, **kwargs)
-        else:
-            raise ValueError(f"Unknown search engine: {engine}")
-
-    def get_name(self) -> str:
-        return BuiltinTool.brave_search.value
-
-    async def run_impl(self, query: str) -> str:
-        return await self.engine.search(query)
-
-
-class BingSearch:
-    def __init__(self, api_key: str, top_k: int = 3, **kwargs) -> None:
-        self.api_key = api_key
-        self.top_k = top_k
-
-    async def search(self, query: str) -> str:
-        url = "https://api.bing.microsoft.com/v7.0/search"
-        headers = {
-            "Ocp-Apim-Subscription-Key": self.api_key,
-        }
-        params = {
-            "count": self.top_k,
-            "textDecorations": True,
-            "textFormat": "HTML",
-            "q": query,
-        }
-
-        response = requests.get(url=url, params=params, headers=headers)
-        response.raise_for_status()
-        clean = self._clean_response(response.json())
-        return json.dumps(clean)
-
-    def _clean_response(self, search_response):
-        clean_response = []
-        query = search_response["queryContext"]["originalQuery"]
-        if "webPages" in search_response:
-            pages = search_response["webPages"]["value"]
-            for p in pages:
-                selected_keys = {"name", "url", "snippet"}
-                clean_response.append(
-                    {k: v for k, v in p.items() if k in selected_keys}
-                )
-        if "news" in search_response:
-            clean_news = []
-            news = search_response["news"]["value"]
-            for n in news:
-                selected_keys = {"name", "url", "description"}
-                clean_news.append({k: v for k, v in n.items() if k in selected_keys})
-
-            clean_response.append(clean_news)
-
-        return {"query": query, "top_k": clean_response}
-
-
-class BraveSearch:
-    def __init__(self, api_key: str) -> None:
-        self.api_key = api_key
-
-    async def search(self, query: str) -> str:
-        url = "https://api.search.brave.com/res/v1/web/search"
-        headers = {
-            "X-Subscription-Token": self.api_key,
-            "Accept-Encoding": "gzip",
-            "Accept": "application/json",
-        }
-        payload = {"q": query}
-        response = requests.get(url=url, params=payload, headers=headers)
-        return json.dumps(self._clean_brave_response(response.json()))
-
-    def _clean_brave_response(self, search_response, top_k=3):
-        query = None
-        clean_response = []
-        if "query" in search_response:
-            if "original" in search_response["query"]:
-                query = search_response["query"]["original"]
-        if "mixed" in search_response:
-            mixed_results = search_response["mixed"]
-            for m in mixed_results["main"][:top_k]:
-                r_type = m["type"]
-                results = search_response[r_type]["results"]
-                if r_type == "web":
-                    # For web data - add a single output from the search
-                    idx = m["index"]
-                    selected_keys = [
-                        "type",
-                        "title",
-                        "url",
-                        "description",
-                        "date",
-                        "extra_snippets",
-                    ]
-                    cleaned = {
-                        k: v for k, v in results[idx].items() if k in selected_keys
-                    }
-                elif r_type == "faq":
-                    # For faw data - take a list of all the questions & answers
-                    selected_keys = ["type", "question", "answer", "title", "url"]
-                    cleaned = []
-                    for q in results:
-                        cleaned.append(
-                            {k: v for k, v in q.items() if k in selected_keys}
-                        )
-                elif r_type == "infobox":
-                    idx = m["index"]
-                    selected_keys = [
-                        "type",
-                        "title",
-                        "url",
-                        "description",
-                        "long_desc",
-                    ]
-                    cleaned = {
-                        k: v for k, v in results[idx].items() if k in selected_keys
-                    }
-                elif r_type == "videos":
-                    selected_keys = [
-                        "type",
-                        "url",
-                        "title",
-                        "description",
-                        "date",
-                    ]
-                    cleaned = []
-                    for q in results:
-                        cleaned.append(
-                            {k: v for k, v in q.items() if k in selected_keys}
-                        )
-                elif r_type == "locations":
-                    # For faw data - take a list of all the questions & answers
-                    selected_keys = [
-                        "type",
-                        "title",
-                        "url",
-                        "description",
-                        "coordinates",
-                        "postal_address",
-                        "contact",
-                        "rating",
-                        "distance",
-                        "zoom_level",
-                    ]
-                    cleaned = []
-                    for q in results:
-                        cleaned.append(
-                            {k: v for k, v in q.items() if k in selected_keys}
-                        )
-                elif r_type == "news":
-                    # For faw data - take a list of all the questions & answers
-                    selected_keys = [
-                        "type",
-                        "title",
-                        "url",
-                        "description",
-                    ]
-                    cleaned = []
-                    for q in results:
-                        cleaned.append(
-                            {k: v for k, v in q.items() if k in selected_keys}
-                        )
-                else:
-                    cleaned = []
-
-                clean_response.append(cleaned)
-
-        return {"query": query, "top_k": clean_response}
-
-
-class WolframAlphaTool(SingleMessageBuiltinTool):
-    def __init__(self, api_key: str) -> None:
-        self.api_key = api_key
-        self.url = "https://api.wolframalpha.com/v2/query"
-
-    def get_name(self) -> str:
-        return BuiltinTool.wolfram_alpha.value
-
-    async def run_impl(self, query: str) -> str:
-        params = {
-            "input": query,
-            "appid": self.api_key,
-            "format": "plaintext",
-            "output": "json",
-        }
-        response = requests.get(
-            self.url,
-            params=params,
-        )
-
-        return json.dumps(self._clean_wolfram_alpha_response(response.json()))
-
-    def _clean_wolfram_alpha_response(self, wa_response):
-        remove = {
-            "queryresult": [
-                "datatypes",
-                "error",
-                "timedout",
-                "timedoutpods",
-                "numpods",
-                "timing",
-                "parsetiming",
-                "parsetimedout",
-                "recalculate",
-                "id",
-                "host",
-                "server",
-                "related",
-                "version",
-                {
-                    "pods": [
-                        "scanner",
-                        "id",
-                        "error",
-                        "expressiontypes",
-                        "states",
-                        "infos",
-                        "position",
-                        "numsubpods",
-                    ]
-                },
-                "assumptions",
-            ],
-        }
-        for main_key in remove:
-            for key_to_remove in remove[main_key]:
-                try:
-                    if key_to_remove == "assumptions":
-                        if "assumptions" in wa_response[main_key]:
-                            del wa_response[main_key][key_to_remove]
-                    if isinstance(key_to_remove, dict):
-                        for sub_key in key_to_remove:
-                            if sub_key == "pods":
-                                for i in range(len(wa_response[main_key][sub_key])):
-                                    if (
-                                        wa_response[main_key][sub_key][i]["title"]
-                                        == "Result"
-                                    ):
-                                        del wa_response[main_key][sub_key][i + 1 :]
-                                        break
-                            sub_items = wa_response[main_key][sub_key]
-                            for i in range(len(sub_items)):
-                                for sub_key_to_remove in key_to_remove[sub_key]:
-                                    if sub_key_to_remove in sub_items[i]:
-                                        del sub_items[i][sub_key_to_remove]
-                    elif key_to_remove in wa_response[main_key]:
-                        del wa_response[main_key][key_to_remove]
-                except KeyError:
-                    pass
-        return wa_response
-
-
-class CodeInterpreterTool(BaseTool):
-    def __init__(self) -> None:
-        ctx = CodeExecutionContext(
-            matplotlib_dump_dir=tempfile.mkdtemp(),
-        )
-        self.code_executor = CodeExecutor(ctx)
-
-    def get_name(self) -> str:
-        return BuiltinTool.code_interpreter.value
-
-    async def run(self, messages: List[CompletionMessage]) -> List[ToolResponseMessage]:
-        message = messages[0]
-        assert len(message.tool_calls) == 1, "Expected a single tool call"
-
-        tool_call = messages[0].tool_calls[0]
-        script = tool_call.arguments["code"]
-
-        req = CodeExecutionRequest(scripts=[script])
-        res = self.code_executor.execute(req)
-
-        pieces = [res["process_status"]]
-        for out_type in ["stdout", "stderr"]:
-            res_out = res[out_type]
-            if res_out != "":
-                pieces.extend([f"[{out_type}]", res_out, f"[/{out_type}]"])
-                if out_type == "stderr":
-                    cprint(f"ipython tool error: ↓\n{res_out}", color="red")
-
-        message = ToolResponseMessage(
-            call_id=tool_call.call_id,
-            tool_name=tool_call.tool_name,
-            content="\n".join(pieces),
-        )
-        return [message]