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

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

@@ -1,161 +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 logging
-import uuid
-from typing import AsyncGenerator
-
-from llama_stack.apis.inference import Inference
-from llama_stack.apis.memory import Memory
-from llama_stack.apis.memory_banks import MemoryBanks
-from llama_stack.apis.safety import Safety
-from llama_stack.apis.agents import *  # noqa: F403
-
-from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
-
-from .agent_instance import ChatAgent
-from .config import MetaReferenceAgentsImplConfig
-
-logger = logging.getLogger()
-logger.setLevel(logging.INFO)
-
-
-class MetaReferenceAgentsImpl(Agents):
-    def __init__(
-        self,
-        config: MetaReferenceAgentsImplConfig,
-        inference_api: Inference,
-        memory_api: Memory,
-        safety_api: Safety,
-        memory_banks_api: MemoryBanks,
-    ):
-        self.config = config
-        self.inference_api = inference_api
-        self.memory_api = memory_api
-        self.safety_api = safety_api
-        self.memory_banks_api = memory_banks_api
-
-        self.in_memory_store = InmemoryKVStoreImpl()
-
-    async def initialize(self) -> None:
-        self.persistence_store = await kvstore_impl(self.config.persistence_store)
-
-    async def create_agent(
-        self,
-        agent_config: AgentConfig,
-    ) -> AgentCreateResponse:
-        agent_id = str(uuid.uuid4())
-
-        await self.persistence_store.set(
-            key=f"agent:{agent_id}",
-            value=agent_config.json(),
-        )
-        return AgentCreateResponse(
-            agent_id=agent_id,
-        )
-
-    async def get_agent(self, agent_id: str) -> ChatAgent:
-        agent_config = await self.persistence_store.get(
-            key=f"agent:{agent_id}",
-        )
-        if not agent_config:
-            raise ValueError(f"Could not find agent config for {agent_id}")
-
-        try:
-            agent_config = json.loads(agent_config)
-        except json.JSONDecodeError as e:
-            raise ValueError(
-                f"Could not JSON decode agent config for {agent_id}"
-            ) from e
-
-        try:
-            agent_config = AgentConfig(**agent_config)
-        except Exception as e:
-            raise ValueError(
-                f"Could not validate(?) agent config for {agent_id}"
-            ) from e
-
-        return ChatAgent(
-            agent_id=agent_id,
-            agent_config=agent_config,
-            inference_api=self.inference_api,
-            safety_api=self.safety_api,
-            memory_api=self.memory_api,
-            memory_banks_api=self.memory_banks_api,
-            persistence_store=(
-                self.persistence_store
-                if agent_config.enable_session_persistence
-                else self.in_memory_store
-            ),
-        )
-
-    async def create_agent_session(
-        self,
-        agent_id: str,
-        session_name: str,
-    ) -> AgentSessionCreateResponse:
-        agent = await self.get_agent(agent_id)
-
-        session_id = await agent.create_session(session_name)
-        return AgentSessionCreateResponse(
-            session_id=session_id,
-        )
-
-    def create_agent_turn(
-        self,
-        agent_id: str,
-        session_id: str,
-        messages: List[
-            Union[
-                UserMessage,
-                ToolResponseMessage,
-            ]
-        ],
-        attachments: Optional[List[Attachment]] = None,
-        stream: Optional[bool] = False,
-    ) -> AsyncGenerator:
-        request = AgentTurnCreateRequest(
-            agent_id=agent_id,
-            session_id=session_id,
-            messages=messages,
-            attachments=attachments,
-            stream=True,
-        )
-        if stream:
-            return self._create_agent_turn_streaming(request)
-        else:
-            raise NotImplementedError("Non-streaming agent turns not yet implemented")
-
-    async def _create_agent_turn_streaming(
-        self,
-        request: AgentTurnCreateRequest,
-    ) -> AsyncGenerator:
-        agent = await self.get_agent(request.agent_id)
-        async for event in agent.create_and_execute_turn(request):
-            yield event
-
-    async def get_agents_turn(self, agent_id: str, turn_id: str) -> Turn:
-        raise NotImplementedError()
-
-    async def get_agents_step(
-        self, agent_id: str, turn_id: str, step_id: str
-    ) -> AgentStepResponse:
-        raise NotImplementedError()
-
-    async def get_agents_session(
-        self,
-        agent_id: str,
-        session_id: str,
-        turn_ids: Optional[List[str]] = None,
-    ) -> Session:
-        raise NotImplementedError()
-
-    async def delete_agents_session(self, agent_id: str, session_id: str) -> None:
-        raise NotImplementedError()
-
-    async def delete_agents(self, agent_id: str) -> None:
-        raise NotImplementedError()

llama_stack/providers/impls/meta_reference/agents/persistence.py

@@ -1,84 +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 uuid
-from datetime import datetime
-
-from typing import List, Optional
-from llama_stack.apis.agents import *  # noqa: F403
-from pydantic import BaseModel
-
-from llama_stack.providers.utils.kvstore import KVStore
-
-
-class AgentSessionInfo(BaseModel):
-    session_id: str
-    session_name: str
-    memory_bank_id: Optional[str] = None
-    started_at: datetime
-
-
-class AgentPersistence:
-    def __init__(self, agent_id: str, kvstore: KVStore):
-        self.agent_id = agent_id
-        self.kvstore = kvstore
-
-    async def create_session(self, name: str) -> str:
-        session_id = str(uuid.uuid4())
-        session_info = AgentSessionInfo(
-            session_id=session_id,
-            session_name=name,
-            started_at=datetime.now(),
-        )
-        await self.kvstore.set(
-            key=f"session:{self.agent_id}:{session_id}",
-            value=session_info.json(),
-        )
-        return session_id
-
-    async def get_session_info(self, session_id: str) -> Optional[AgentSessionInfo]:
-        value = await self.kvstore.get(
-            key=f"session:{self.agent_id}:{session_id}",
-        )
-        if not value:
-            return None
-
-        return AgentSessionInfo(**json.loads(value))
-
-    async def add_memory_bank_to_session(self, session_id: str, bank_id: str):
-        session_info = await self.get_session_info(session_id)
-        if session_info is None:
-            raise ValueError(f"Session {session_id} not found")
-
-        session_info.memory_bank_id = bank_id
-        await self.kvstore.set(
-            key=f"session:{self.agent_id}:{session_id}",
-            value=session_info.json(),
-        )
-
-    async def add_turn_to_session(self, session_id: str, turn: Turn):
-        await self.kvstore.set(
-            key=f"session:{self.agent_id}:{session_id}:{turn.turn_id}",
-            value=turn.json(),
-        )
-
-    async def get_session_turns(self, session_id: str) -> List[Turn]:
-        values = await self.kvstore.range(
-            start_key=f"session:{self.agent_id}:{session_id}:",
-            end_key=f"session:{self.agent_id}:{session_id}:\xff\xff\xff\xff",
-        )
-        turns = []
-        for value in values:
-            try:
-                turn = Turn(**json.loads(value))
-                turns.append(turn)
-            except Exception as e:
-                print(f"Error parsing turn: {e}")
-                continue
-
-        return turns

llama_stack/providers/impls/meta_reference/agents/rag/context_retriever.py

@@ -1,74 +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 List
-
-from jinja2 import Template
-from llama_models.llama3.api import *  # noqa: F403
-
-
-from termcolor import cprint  # noqa: F401
-
-from llama_stack.apis.agents import (
-    DefaultMemoryQueryGeneratorConfig,
-    LLMMemoryQueryGeneratorConfig,
-    MemoryQueryGenerator,
-    MemoryQueryGeneratorConfig,
-)
-from llama_stack.apis.inference 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 = await inference_api.chat_completion(
-        model=model,
-        messages=[message],
-        stream=False,
-    )
-
-    query = response.completion_message.content
-
-    return query

llama_stack/providers/impls/meta_reference/agents/safety.py

@@ -1,57 +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 asyncio
-
-from typing import List
-
-from llama_models.llama3.api.datatypes import Message
-from termcolor import cprint
-
-from llama_stack.apis.safety import *  # noqa: F403
-
-
-class SafetyException(Exception):  # noqa: N818
-    def __init__(self, violation: SafetyViolation):
-        self.violation = violation
-        super().__init__(violation.user_message)
-
-
-class ShieldRunnerMixin:
-    def __init__(
-        self,
-        safety_api: Safety,
-        input_shields: List[str] = None,
-        output_shields: List[str] = None,
-    ):
-        self.safety_api = safety_api
-        self.input_shields = input_shields
-        self.output_shields = output_shields
-
-    async def run_multiple_shields(
-        self, messages: List[Message], shield_types: List[str]
-    ) -> None:
-        responses = await asyncio.gather(
-            *[
-                self.safety_api.run_shield(
-                    shield_type=shield_type,
-                    messages=messages,
-                )
-                for shield_type in shield_types
-            ]
-        )
-        for shield_type, response in zip(shield_types, responses):
-            if not response.violation:
-                continue
-
-            violation = response.violation
-            if violation.violation_level == ViolationLevel.ERROR:
-                raise SafetyException(violation)
-            elif violation.violation_level == ViolationLevel.WARN:
-                cprint(
-                    f"[Warn]{shield_type} raised a warning",
-                    color="red",
-                )

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

@@ -1,93 +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 unittest
-
-from llama_models.llama3.api.datatypes import (
-    Attachment,
-    BuiltinTool,
-    CompletionMessage,
-    StopReason,
-    ToolCall,
-)
-
-from ..tools.builtin import CodeInterpreterTool
-
-
-class TestCodeInterpreter(unittest.IsolatedAsyncioTestCase):
-    async def test_matplotlib(self):
-        tool = CodeInterpreterTool()
-        code = """
-import matplotlib.pyplot as plt
-import numpy as np
-
-x = np.array([1, 1])
-y = np.array([0, 10])
-
-plt.plot(x, y)
-plt.title('x = 1')
-plt.xlabel('x')
-plt.ylabel('y')
-plt.grid(True)
-plt.axvline(x=1, color='r')
-plt.show()
-        """
-        message = CompletionMessage(
-            role="assistant",
-            content="",
-            tool_calls=[
-                ToolCall(
-                    call_id="call_id",
-                    tool_name=BuiltinTool.code_interpreter,
-                    arguments={"code": code},
-                )
-            ],
-            stop_reason=StopReason.end_of_message,
-        )
-        ret = await tool.run([message])
-
-        self.assertEqual(len(ret), 1)
-
-        output = ret[0].content
-        self.assertIsInstance(output, Attachment)
-        self.assertEqual(output.mime_type, "image/png")
-
-    async def test_path_unlink(self):
-        tool = CodeInterpreterTool()
-        code = """
-import os
-from pathlib import Path
-import tempfile
-
-dpath = Path(os.environ["MPLCONFIGDIR"])
-with open(dpath / "test", "w") as f:
-    f.write("hello")
-
-Path(dpath / "test").unlink()
-print("_OK_")
-        """
-        message = CompletionMessage(
-            role="assistant",
-            content="",
-            tool_calls=[
-                ToolCall(
-                    call_id="call_id",
-                    tool_name=BuiltinTool.code_interpreter,
-                    arguments={"code": code},
-                )
-            ],
-            stop_reason=StopReason.end_of_message,
-        )
-        ret = await tool.run([message])
-
-        self.assertEqual(len(ret), 1)
-
-        output = ret[0].content
-        self.assertTrue("_OK_" in output)
-
-
-if __name__ == "__main__":
-    unittest.main()