openbb-agent 0.1.0__tar.gz

openbb_agent-0.1.0/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: openbb-agent
+Version: 0.1.0
+Summary: OpenBB Terminal Pro copilot agent
+Author-email: hjm_g <ghj791990@gmail.com>
+License-Expression: MIT
+Project-URL: Repository, https://github.com/Findworth/openbb_agent
+Keywords: openbb,finance,agent,copilot
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Requires-Dist: openbb-ai
+Requires-Dist: fastapi
+Requires-Dist: uvicorn[standard]
+Requires-Dist: sse-starlette
+Requires-Dist: python-dotenv
+Requires-Dist: pydantic-settings
+Requires-Dist: magentic
+Requires-Dist: crewai
+Requires-Dist: redis
+Requires-Dist: mem0ai
+Requires-Dist: chromadb
+Provides-Extra: dev
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: pytest; extra == "dev"

openbb_agent-0.1.0/openbb_agent/agent.py

@@ -0,0 +1,440 @@
+import asyncio
+import inspect
+import json
+from typing import AsyncGenerator
+
+from magentic import (
+    AsyncStreamedResponse,
+    Chat,
+    AsyncStreamedStr,
+    FunctionCall,
+    OpenaiChatModel,
+)
+from .prompts import PromptManager, render_system_prompt
+from .memory.summarizer import summarize_and_store
+from .models import (
+    ClientFunctionCallError,
+    LlmClientMessage,
+    QueryRequest,
+    Citation,
+    CitationCollection,
+    CitationCollectionSSE,
+    DataContent,
+    DataFileReferences,
+    DataSourceRequest,
+    FunctionCallSSE,
+    FunctionCallSSEData,
+    LlmClientFunctionCallResultMessage,
+    LlmFunctionCall,
+    MessageChunkSSE,
+    MessageChunkSSEData,
+    StatusUpdateSSE,
+    StatusUpdateSSEData,
+    Widget,
+)
+from magentic import (
+    AssistantMessage,
+    FunctionResultMessage,
+    AnyMessage,
+    SystemMessage,
+    UserMessage,
+)
+from typing import Any, AsyncGenerator, Awaitable, Callable, Literal, Protocol, cast
+from openai import AsyncOpenAI
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionUserMessageParam,
+    ChatCompletionAssistantMessageParam,
+)
+
+
+def reasoning_step(
+    event_type: Literal["INFO", "WARNING", "ERROR"],
+    message: str,
+    details: dict[str, Any] | None = None,
+) -> StatusUpdateSSE:
+    return StatusUpdateSSE(
+        data=StatusUpdateSSEData(
+            eventType=event_type,
+            message=message,
+            details=[details] if details else [],
+        )
+    )
+
+
+class WrappedFunctionProtocol(Protocol):
+    async def execute_post_processing(
+        self, data: list[DataContent | DataFileReferences | ClientFunctionCallError]
+    ) -> str: ...
+
+    def execute_callbacks(
+        self,
+        function_call_result: LlmClientFunctionCallResultMessage,
+        request: QueryRequest,
+    ) -> AsyncGenerator[Any, None]: ...
+
+    def __call__(
+        self, *args: Any, **kwargs: Any
+    ) -> AsyncGenerator[FunctionCallSSE | StatusUpdateSSE, None]: ...
+
+
+def remote_function_call(
+    function: Literal["get_widget_data"],
+    output_formatter: Callable[..., Awaitable[str]] | None = None,
+    callbacks: list[Callable[..., Awaitable[Any]]] | None = None,
+) -> Callable:
+    if function not in ["get_widget_data"]:
+        raise ValueError(f"Unsupported function: {function}. Must be 'get_widget_data'.")
+
+    def outer_wrapper(func: Callable) -> WrappedFunctionProtocol:
+        class InnerWrapper(WrappedFunctionProtocol):
+            def __init__(self):
+                self.__name__ = func.__name__
+                self.__signature__ = self._mask_signature(func)
+                self.__doc__ = func.__doc__
+                self.local_function = func
+                self.function = function
+                self.post_process_function = output_formatter
+                self.callbacks = callbacks
+                self._request = None
+
+            @property
+            def request(self) -> QueryRequest:
+                return self._request
+
+            @request.setter
+            def request(self, request: QueryRequest):
+                self._request = request
+
+            def _mask_signature(self, func: Callable):
+                signature = inspect.signature(func)
+                masked_params = [p for p in signature.parameters.values() if p.name != "request"]
+                return signature.replace(parameters=masked_params)
+
+            async def execute_callbacks(
+                self,
+                function_call_result: LlmClientFunctionCallResultMessage,
+                request: QueryRequest,
+            ) -> AsyncGenerator[Any, None]:
+                if self.callbacks:
+                    for callback in self.callbacks:
+                        if inspect.isasyncgenfunction(callback):
+                            async for event in callback(function_call_result, request):
+                                yield event
+                        else:
+                            await callback(function_call_result, self.request)
+
+            async def execute_post_processing(
+                self,
+                data: list[DataContent | DataFileReferences | ClientFunctionCallError],
+            ) -> str:
+                if self.post_process_function:
+                    return await self.post_process_function(data)
+                return str(data)
+
+            async def __call__(
+                self, *args, **kwargs
+            ) -> AsyncGenerator[FunctionCallSSE | StatusUpdateSSE, None]:
+                bound_args = self.__signature__.bind(*args, **kwargs).arguments
+                async for event in func(*args, request=self._request, **kwargs):
+                    if isinstance(event, StatusUpdateSSE):
+                        yield event
+                    elif isinstance(event, DataSourceRequest):
+                        yield FunctionCallSSE(
+                            data=FunctionCallSSEData(
+                                function=self.function,
+                                input_arguments={
+                                    "data_sources": [
+                                        DataSourceRequest(
+                                            widget_uuid=event.widget_uuid,
+                                            origin=event.origin,
+                                            id=event.id,
+                                            input_args=event.input_args,
+                                        )
+                                    ]
+                                },
+                                extra_state={
+                                    "copilot_function_call_arguments": {**bound_args},
+                                    "_locally_bound_function": func.__name__,
+                                },
+                            )
+                        )
+                        return
+                    else:
+                        yield event
+
+        return InnerWrapper()
+
+    return outer_wrapper
+
+
+def get_remote_data(
+    widget: Widget,
+    input_arguments: dict[str, Any],
+) -> DataSourceRequest:
+    return DataSourceRequest(
+        widget_uuid=str(widget.uuid),
+        origin=widget.origin,
+        id=widget.widget_id,
+        input_args=input_arguments,
+    )
+
+
+def _last_user_message(request: QueryRequest) -> str:
+    for m in reversed[LlmClientFunctionCallResultMessage | LlmClientMessage](request.messages):
+        if isinstance(m, LlmClientMessage) and m.role == "human":
+            return m.content if isinstance(m.content, str) else ""
+    return ""
+
+
+async def _get_memory_context(
+    memory_service,
+    request: QueryRequest,
+    memory_ids: dict[str, str | None],
+) -> str:
+    if not memory_service or not memory_service.is_available():
+        return ""
+    query = _last_user_message(request)
+    if not query:
+        return ""
+    kwargs: dict = {
+        "query": query,
+        "user_id": memory_ids.get("user_id") or "default",
+        "limit": 5,
+    }
+    if memory_ids.get("agent_id"):
+        kwargs["agent_id"] = memory_ids["agent_id"]
+    if memory_ids.get("app_id"):
+        kwargs["app_id"] = memory_ids["app_id"]
+    results = await asyncio.to_thread(memory_service.search, **kwargs)
+    return memory_service.format_context(results)
+
+
+async def run_copilot(
+    request: QueryRequest,
+    memory_ids: dict[str, str | None],
+    functions: list,
+    model,
+    memory_service=None,
+) -> AsyncGenerator[dict, None]:
+    memory_context = await _get_memory_context(memory_service, request, memory_ids)
+    workspace_options = getattr(request, "workspace_options", []) or []
+    run_crew_enabled = "run-crew" in workspace_options
+    run_flow_enabled = "run-flow" in workspace_options
+    prompt_manager = PromptManager(
+        include_widgets=True,
+        include_flows_crews=run_crew_enabled or run_flow_enabled,
+    )
+    system_prompt = render_system_prompt(
+        request=request,
+        widget_collection=request.widgets,
+        memory_context=memory_context,
+        prompt_manager=prompt_manager,
+    )
+    agent = OpenBBAgent(
+        query_request=request,
+        system_prompt=system_prompt,
+        functions=functions,
+        model=model,
+        prompt_manager=prompt_manager,
+    )
+    async for event in agent.run():
+        yield event
+
+    if memory_service:
+        await summarize_and_store(
+            memory_service,
+            request,
+            memory_ids=memory_ids,
+            n_rounds=3,
+        )
+
+
+def get_wrapped_function(
+    function_name: str, functions: list[Any]
+) -> WrappedFunctionProtocol:
+    matching_local_functions = list(filter(lambda x: x.__name__ == function_name, functions))
+    wrapped_function = matching_local_functions[0] if matching_local_functions else None
+    if not wrapped_function:
+        raise ValueError(f"Local function not found: {function_name}")
+    return wrapped_function
+
+
+class OpenBBAgent:
+    def __init__(
+        self,
+        query_request: QueryRequest,
+        system_prompt: str,
+        functions: list[Callable] | None = None,
+        chat_class: type[Chat] | None = None,
+        model: str | OpenaiChatModel | None = None,
+        hooks: Any = None,
+        prompt_manager: PromptManager | None = None,
+        **kwargs: Any,
+    ):
+        self.request = query_request
+        self.widgets = query_request.widgets
+        self.system_prompt = system_prompt
+        self.functions = functions
+        self.chat_class = chat_class or Chat
+        self._model = model if model is not None else OpenaiChatModel()
+        self._chat: Chat | None = None
+        self._citations: CitationCollection | None = None
+        self._messages: list[AnyMessage] = []
+        self._hooks = hooks
+        self._kwargs = kwargs
+        self.prompt_manager = prompt_manager or PromptManager()
+
+    async def run(self, max_completions: int = 10) -> AsyncGenerator[dict, None]:
+        self._messages = await self._handle_request()
+        self._citations = await self._handle_callbacks()
+        self._chat = self.chat_class(
+            messages=self._messages,
+            output_types=[AsyncStreamedResponse],
+            functions=self.functions if self.functions else None,
+            model=self._model,  # type: ignore[arg-type]
+            **self._kwargs,
+        )
+        async for event in self._execute(max_completions=max_completions):
+            yield event.model_dump()
+        if self._citations.citations:
+            yield CitationCollectionSSE(data=self._citations).model_dump()
+
+    async def _handle_callbacks(self) -> CitationCollection:
+        if not self.functions:
+            return CitationCollection(citations=[])
+        citations: list[Citation] = []
+        if isinstance(self.request.messages[-1], LlmClientFunctionCallResultMessage):
+            wrapped_function = get_wrapped_function(
+                function_name=self.request.messages[-1].function,
+                functions=self.functions,
+            )
+            if hasattr(wrapped_function, "execute_callbacks"):
+                async for event in wrapped_function.execute_callbacks(  # type: ignore
+                    request=self.request, function_call_result=self.request.messages[-1]
+                ):
+                    if isinstance(event, Citation):
+                        citations.append(event)
+        return CitationCollection(citations=citations)
+
+    async def _handle_request(self) -> list[AnyMessage]:
+        chat_messages: list[AnyMessage] = [SystemMessage(self.system_prompt)]
+        for message in self.request.messages:
+            match message:
+                case LlmClientMessage(role="human"):
+                    chat_messages.append(UserMessage(content=message.content))
+                case LlmClientMessage(role="ai") if isinstance(message.content, str):
+                    chat_messages.append(AssistantMessage(content=message.content))
+                case LlmClientMessage(role="ai") if isinstance(message.content, LlmFunctionCall):
+                    pass
+                case LlmClientFunctionCallResultMessage(role="tool"):
+                    if not self.functions:
+                        continue
+                    fn_name = message.extra_state.get("_locally_bound_function") or message.function
+                    wrapped_function = get_wrapped_function(
+                        function_name=fn_name,
+                        functions=self.functions,
+                    )
+                    function_call = FunctionCall(
+                        function=wrapped_function,
+                        **message.extra_state.get("copilot_function_call_arguments", {}),
+                    )
+                    chat_messages.append(AssistantMessage(function_call))
+                    if hasattr(wrapped_function, "execute_post_processing"):
+                        content = await wrapped_function.execute_post_processing(message.data)
+                    else:
+                        content = str(message.data)
+                    chat_messages.append(
+                        FunctionResultMessage(
+                            content=content,
+                            function_call=function_call,
+                        )
+                    )
+                case _:
+                    raise ValueError(f"Unsupported message type: {message}")
+        return chat_messages
+
+    async def _handle_text_stream(
+        self, stream: AsyncStreamedStr
+    ) -> AsyncGenerator[MessageChunkSSE, None]:
+        self._chat = cast(Chat, self._chat)
+        async for chunk in stream:
+            yield MessageChunkSSE(data=MessageChunkSSEData(delta=chunk))
+
+    async def _handle_function_call(
+        self, function_call: FunctionCall
+    ) -> AsyncGenerator[FunctionCallSSE | StatusUpdateSSE | Any, None]:
+        self._chat = cast(Chat, self._chat)
+        function_call_result: str = ""
+        if not isinstance(self._chat.last_message, AssistantMessage):
+            raise ValueError("Last message is not an assistant message")
+        if hasattr(function_call.function, "request"):
+            function_call.function.request = self.request
+        if self._hooks is not None and hasattr(self._hooks, "run"):
+            await self._hooks.run("pre_tool_call", function_call=function_call)
+        yield reasoning_step(
+            "INFO",
+            "Tool execution started",
+            {"function": getattr(function_call.function, "__name__", "unknown")},
+        )
+        async for event in function_call():
+            if isinstance(event, StatusUpdateSSE):
+                yield event
+            elif isinstance(event, FunctionCallSSE):
+                yield event
+                return
+            elif hasattr(event, "event") and getattr(event, "event", None) == "copilotMessageArtifact":
+                yield event
+            else:
+                function_call_result += str(event)
+        yield reasoning_step(
+            "INFO",
+            "Tool execution completed",
+            {"function": getattr(function_call.function, "__name__", "unknown")},
+        )
+        if self._hooks is not None and hasattr(self._hooks, "run"):
+            await self._hooks.run(
+                "post_tool_call",
+                function_call=function_call,
+                result=function_call_result,
+            )
+        self._chat = self._chat.add_message(
+            FunctionResultMessage(
+                content=function_call_result,
+                function_call=function_call,
+            )
+        )
+
+    async def _execute(
+        self, max_completions: int
+    ) -> AsyncGenerator[MessageChunkSSE | FunctionCallSSE | StatusUpdateSSE, None]:
+        completion_count = 0
+        while completion_count < max_completions:
+            completion_count += 1
+            yield reasoning_step(
+                "INFO",
+                f"LLM reasoning (round {completion_count})...",
+                {"completion_count": completion_count},
+            )
+            if self._hooks is not None and hasattr(self._hooks, "run"):
+                await self._hooks.run("pre_llm", request=self.request)
+            self._chat = await cast(Chat, self._chat).asubmit()
+            if self._hooks is not None and hasattr(self._hooks, "run"):
+                await self._hooks.run("post_llm", chat=self._chat, request=self.request)
+            if isinstance(self._chat.last_message.content, AsyncStreamedResponse):
+                async for item in self._chat.last_message.content:
+                    if isinstance(item, AsyncStreamedStr):
+                        async for event in self._handle_text_stream(item):
+                            yield event
+                    elif isinstance(item, FunctionCall):
+                        yield reasoning_step(
+                            "INFO",
+                            "Executing tool/function call...",
+                            {"function": getattr(item.function, "__name__", str(item.function))},
+                        )
+                        async for event in self._handle_function_call(item):
+                            yield event
+                            if isinstance(event, FunctionCallSSE):
+                                return
+                return

openbb_agent-0.1.0/openbb_agent/cache.py

@@ -0,0 +1,74 @@
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+_env_cache_path = Path(__file__).resolve().parent.parent / ".env.cache"
+
+
+class FunctionCacheConfig(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_file=_env_cache_path,
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+    redis_url: str = ""
+    ttl_seconds: int = 3600
+
+
+class FunctionCache:
+    def __init__(self, config: FunctionCacheConfig | None = None) -> None:
+        cfg = config or FunctionCacheConfig()
+        self._redis_url = cfg.redis_url
+        self._ttl = cfg.ttl_seconds
+        self._client: Any = None
+
+    def _get_client(self) -> Any:
+        if self._client is not None:
+            return self._client
+        if not self._redis_url:
+            return None
+        try:
+            import redis.asyncio as redis
+
+            self._client = redis.from_url(self._redis_url, decode_responses=True)
+            return self._client
+        except ImportError:
+            return None
+
+    def _cache_key(self, func_name: str, args: dict[str, Any]) -> str:
+        payload = json.dumps(sorted(args.items()), sort_keys=True, default=str)
+        h = hashlib.sha256(f"{func_name}:{payload}".encode()).hexdigest()
+        return f"fc:{func_name}:{h}"
+
+    async def get(self, func_name: str, args: dict[str, Any]) -> str | None:
+        client = self._get_client()
+        if client is None:
+            return None
+        try:
+            key = self._cache_key(func_name, args)
+            return await client.get(key)
+        except Exception:
+            return None
+
+    async def set(
+        self,
+        func_name: str,
+        args: dict[str, Any],
+        value: str,
+        ttl: int | None = None,
+    ) -> None:
+        client = self._get_client()
+        if client is None:
+            return
+        try:
+            key = self._cache_key(func_name, args)
+            await client.set(key, value, ex=ttl or self._ttl)
+        except Exception:
+            pass
+
+    def is_enabled(self) -> bool:
+        return bool(self._redis_url)

openbb_agent-0.1.0/openbb_agent/config.py

@@ -0,0 +1,49 @@
+from functools import lru_cache
+from typing import Any
+
+from magentic import OpenaiChatModel
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class ModelConfig(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+    openai_api_key: str = ""
+    openai_base_url: str | None = None
+    deepseek_api_key: str = ""
+    deepseek_base_url: str = "https://api.deepseek.com"
+    model_name: str = "deepseek-chat"
+    delegate_model_name: str | None = None
+
+
+def get_chat_model(
+    api_key: str | None = None,
+    base_url: str | None = None,
+    model: str | None = None,
+) -> OpenaiChatModel:
+    cfg = get_config()
+    model_name = model or cfg.model_name
+    if "deepseek" in model_name.lower():
+        key = api_key or cfg.deepseek_api_key or None
+        url = base_url or cfg.deepseek_base_url or None
+    else:
+        key = api_key or cfg.openai_api_key or None
+        url = base_url or cfg.openai_base_url or None
+    return OpenaiChatModel(
+        model=model_name,
+        api_key=key,
+        base_url=url,
+    )
+
+
+@lru_cache
+def _get_config() -> ModelConfig:
+    return ModelConfig()
+
+
+def get_config() -> ModelConfig:
+    return _get_config()

openbb_agent-0.1.0/openbb_agent/data/disabled_bash_commands.txt

@@ -0,0 +1,8 @@
+rm
+sudo
+mkfs
+dd
+chmod
+chown
+wget
+curl

openbb_agent-0.1.0/openbb_agent/hooks.py

@@ -0,0 +1,47 @@
+from enum import Enum
+from typing import Any, Awaitable, Callable
+
+
+class HookType(str, Enum):
+    PRE_LLM = "pre_llm"
+    POST_LLM = "post_llm"
+    PRE_TOOL_CALL = "pre_tool_call"
+    POST_TOOL_CALL = "post_tool_call"
+
+
+HookFunc = Callable[..., Awaitable[None]]
+
+
+class HooksRegistry:
+    def __init__(self) -> None:
+        self._hooks: dict[HookType, list[HookFunc]] = {
+            t: [] for t in HookType
+        }
+
+    def register(self, hook_type: HookType) -> Callable[[HookFunc], HookFunc]:
+        def decorator(fn: HookFunc) -> HookFunc:
+            self._hooks[hook_type].append(fn)
+            return fn
+        return decorator
+
+    async def run(
+        self,
+        hook_type: HookType | str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        ht = HookType(hook_type) if isinstance(hook_type, str) else hook_type
+        for fn in self._hooks.get(ht, []):
+            await fn(*args, **kwargs)
+
+
+def register_hook(hook_type: HookType, registry: HooksRegistry | None = None) -> Callable[[HookFunc], HookFunc]:
+    _registry = registry or _default_registry
+
+    def decorator(fn: HookFunc) -> HookFunc:
+        _registry._hooks[hook_type].append(fn)
+        return fn
+    return decorator
+
+
+_default_registry = HooksRegistry()