From 92c10fc41ebe36455e2258f995c75dd77936e7e0 Mon Sep 17 00:00:00 2001 From: Siddhesh Agarwal Date: Fri, 1 Nov 2024 11:07:04 +0530 Subject: [PATCH] added logging --- simplemind/logging.py | 28 ++++++++++++++++++++++++++++ simplemind/providers/anthropic.py | 4 ++++ simplemind/providers/gemini.py | 4 ++++ simplemind/providers/groq.py | 4 ++++ simplemind/providers/ollama.py | 4 ++++ simplemind/providers/openai.py | 10 +++++++--- simplemind/providers/xai.py | 12 +++++++++++- 7 files changed, 62 insertions(+), 4 deletions(-) create mode 100644 simplemind/logging.py diff --git a/simplemind/logging.py b/simplemind/logging.py new file mode 100644 index 0000000..6dcf779 --- /dev/null +++ b/simplemind/logging.py @@ -0,0 +1,28 @@ +import time +from typing import Any, Callable + +import logfire + +from .settings import settings + + +def logger(func: Callable[..., Any]) -> Callable[..., Any]: + """A @logger decorator that logs the function parameters, function returns, and exceptions raised if logging is enabled.""" + is_logging_enabled = settings.logging.enabled + + def wrapper(*args, **kwargs) -> Any: + if not is_logging_enabled: + return func(*args, **kwargs) + logfire.info(f"Calling {func.__name__} with args: {args}, kwargs: {kwargs}") + t1 = time.perf_counter() + try: + result = func(*args, **kwargs) + t2 = time.perf_counter() + logfire.info(f"{func.__name__} returned: {result} in {t2-t1} seconds") + return result + except Exception as e: + t2 = time.perf_counter() + logfire.error(f"Error in {func.__name__}: {e} in {t2-t1} seconds") + raise e + + return wrapper diff --git a/simplemind/providers/anthropic.py b/simplemind/providers/anthropic.py index 5c76160..9978a17 100644 --- a/simplemind/providers/anthropic.py +++ b/simplemind/providers/anthropic.py @@ -5,6 +5,7 @@ import anthropic import instructor from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider @@ -37,6 +38,7 @@ class Anthropic(BaseProvider): """A client patched with Instructor.""" return instructor.from_anthropic(self.client) + @logger def send_conversation(self, conversation: "Conversation", **kwargs): """Send a conversation to the Anthropic API.""" from ..models import Message @@ -63,6 +65,7 @@ class Anthropic(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response( self, response_model: Type[T], *, llm_model: str | None = None, **kwargs ) -> T: @@ -82,6 +85,7 @@ class Anthropic(BaseProvider): ) return response + @logger def generate_text(self, prompt: str, *, llm_model: str, **kwargs): messages = [ {"role": "user", "content": prompt}, diff --git a/simplemind/providers/gemini.py b/simplemind/providers/gemini.py index c902141..4a20262 100644 --- a/simplemind/providers/gemini.py +++ b/simplemind/providers/gemini.py @@ -8,6 +8,7 @@ import google.generativeai as genai import instructor from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider @@ -38,6 +39,7 @@ class Gemini(BaseProvider): """A Gemini client patched with Instructor.""" return instructor.from_gemini(self.client) + @logger def send_conversation(self, conversation: "Conversation") -> "Message": """Send a conversation to the Gemini API.""" from ..models import Message @@ -64,6 +66,7 @@ class Gemini(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response(self, prompt: str, response_model: Type[T], **kwargs) -> T: """Send a structured response to the Gemini API.""" llm_model = kwargs.pop("llm_model", self.model_name) @@ -81,6 +84,7 @@ class Gemini(BaseProvider): ) from e return response + @logger def generate_text(self, prompt: str, **kwargs) -> str: """Generate text using the Gemini API.""" llm_model = kwargs.pop("llm_model", self.model_name) diff --git a/simplemind/providers/groq.py b/simplemind/providers/groq.py index 5a033be..c3a82f8 100644 --- a/simplemind/providers/groq.py +++ b/simplemind/providers/groq.py @@ -5,6 +5,7 @@ import groq import instructor from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider @@ -36,6 +37,7 @@ class Groq(BaseProvider): """A client patched with Instructor.""" return instructor.from_groq(self.client) + @logger def send_conversation( self, conversation: "Conversation", @@ -66,6 +68,7 @@ class Groq(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response(self, prompt: str, response_model: Type[T], **kwargs) -> T: # Ensure messages are provided in kwargs messages = [ @@ -80,6 +83,7 @@ class Groq(BaseProvider): ) return response + @logger def generate_text( self, prompt: str, diff --git a/simplemind/providers/ollama.py b/simplemind/providers/ollama.py index 72e000d..e9c3220 100644 --- a/simplemind/providers/ollama.py +++ b/simplemind/providers/ollama.py @@ -6,6 +6,7 @@ import ollama as ol from openai import OpenAI from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider @@ -45,6 +46,7 @@ class Ollama(BaseProvider): mode=instructor.Mode.JSON, ) + @logger def send_conversation(self, conversation: "Conversation", **kwargs) -> "Message": """Send a conversation to the Ollama API.""" from ..models import Message @@ -68,6 +70,7 @@ class Ollama(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response( self, prompt: str, @@ -89,6 +92,7 @@ class Ollama(BaseProvider): ) return response + @logger def generate_text( self, prompt: str, *, llm_model: str | None = None, **kwargs ) -> str: diff --git a/simplemind/providers/openai.py b/simplemind/providers/openai.py index a756835..f99e9b7 100644 --- a/simplemind/providers/openai.py +++ b/simplemind/providers/openai.py @@ -5,6 +5,7 @@ import instructor import openai as oa from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider @@ -36,6 +37,7 @@ class OpenAI(BaseProvider): """A OpenAI client with Instructor.""" return instructor.from_openai(self.client) + @logger def send_conversation(self, conversation: "Conversation", **kwargs): """Send a conversation to the OpenAI API.""" from ..models import Message @@ -47,7 +49,7 @@ class OpenAI(BaseProvider): response = self.client.chat.completions.create( model=conversation.llm_model or DEFAULT_MODEL, messages=messages, - **{**self.DEFAULT_KWARGS, **kwargs} + **{**self.DEFAULT_KWARGS, **kwargs}, ) # Get the response content from the OpenAI response @@ -62,6 +64,7 @@ class OpenAI(BaseProvider): llm_provider=PROVIDER_NAME, ) + @logger def structured_response( self, prompt: str, @@ -79,10 +82,11 @@ class OpenAI(BaseProvider): messages=messages, model=llm_model or self.DEFAULT_MODEL, response_model=response_model, - **{**self.DEFAULT_KWARGS, **kwargs} + **{**self.DEFAULT_KWARGS, **kwargs}, ) return response + @logger def generate_text(self, prompt: str, *, llm_model: str | None = None, **kwargs): """Generate text using the OpenAI API.""" messages = [ @@ -91,6 +95,6 @@ class OpenAI(BaseProvider): response = self.client.chat.completions.create( messages=messages, model=llm_model or self.DEFAULT_MODEL, - **{**self.DEFAULT_KWARGS, **kwargs} + **{**self.DEFAULT_KWARGS, **kwargs}, ) return response.choices[0].message.content diff --git a/simplemind/providers/xai.py b/simplemind/providers/xai.py index 28ea14d..936bca2 100644 --- a/simplemind/providers/xai.py +++ b/simplemind/providers/xai.py @@ -1,8 +1,11 @@ from functools import cached_property +from typing import Type, TypeVar import instructor import openai as oa +from pydantic import BaseModel +from ..logging import logger from ..settings import settings from ._base import BaseProvider @@ -12,6 +15,8 @@ BASE_URL = "https://api.x.ai/v1" DEFAULT_MAX_TOKENS = 1000 DEFAULT_KWARGS = {"max_tokens": DEFAULT_MAX_TOKENS} +T = TypeVar("T", bound=BaseModel) + class XAI(BaseProvider): NAME = PROVIDER_NAME @@ -36,6 +41,7 @@ class XAI(BaseProvider): """A client patched with Instructor.""" return instructor.from_openai(self.client) + @logger def send_conversation(self, conversation: "Conversation", **kwargs): """Send a conversation to the OpenAI API.""" from ..models import Message @@ -62,9 +68,13 @@ class XAI(BaseProvider): llm_provider=PROVIDER_NAME, ) - def structured_response(self, prompt: str, response_model, *, llm_model: str): + @logger + def structured_response( + self, prompt: str, response_model: Type[T], *, llm_model: str + ) -> T: raise NotImplementedError("XAI does not support structured responses") + @logger def generate_text(self, prompt: str, *, llm_model: str, **kwargs): messages = [ {"role": "user", "content": prompt},