wbcore 1.46.0__py2.py3-none-any.whl → 1.58.2__py2.py3-none-any.whl

wbcore/contrib/ai/exceptions.py

@@ -1,10 +1,12 @@
 from langchain_core.exceptions import LangChainException, OutputParserException
 
 APIStatusErrors = [LangChainException, OutputParserException]
+BadRequestErrors = []
 
 try:
     from openai._exceptions import (
         AuthenticationError,
+        BadRequestError,
         ConflictError,
         InternalServerError,
         NotFoundError,
@@ -24,17 +26,15 @@ try:
             InternalServerError,
         ]
     )
+    BadRequestErrors.append(BadRequestError)
 except ImportError:
     pass
 
 
 try:
-    from anthropic._exceptions import (
-        APIConnectionError,
-        APIResponseValidationError,
-        APIStatusError,
-    )
+    from anthropic._exceptions import APIConnectionError, APIResponseValidationError, APIStatusError, BadRequestError
 
     APIStatusErrors.extend([APIResponseValidationError, APIStatusError, APIConnectionError])
+    BadRequestErrors.append(BadRequestError)
 except ImportError:
     pass

wbcore/contrib/ai/llm/config.py

@@ -1,22 +1,22 @@
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar
+from typing import Any, Callable, Generic, TypeVar
 
 from celery import shared_task
 from django.db import models
 from django.db.models.signals import ModelSignal
 from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import BaseMessage
 from langchain_openai import ChatOpenAI
+from pydantic import BaseModel
 
-from ..exceptions import APIStatusErrors
-from .utils import construct_prompt, run_llm
+from ..exceptions import APIStatusErrors, BadRequestErrors
+from .utils import run_llm
 
 logger = logging.getLogger("llm")
 
-if TYPE_CHECKING:
-    from pydantic import BaseModel
-
 
 @shared_task(
+    queue="llm",
     autoretry_for=tuple(APIStatusErrors),
     retry_backoff=10,
     max_retries=5,  # retry 5 times maximum
@@ -31,16 +31,31 @@ def invoke_as_task(
     max_tokens: int,
     model_field: str | None,
     output_model: "type[BaseModel] | None",
-    llm_kwargs: dict[str, Any],
+    tools: list[dict[str, Any]],
+    query: Callable | dict[str, Any] | None,
+    extra_query: dict[str, Any] | None,
+    llm_kwargs: dict[str, Any] | None,
+    result_parser: Callable | None,
 ):
     try:
-        result = run_llm(prompt, instance, output_model, chat_model, chat_model_name, max_tokens, **llm_kwargs)
-        if model_field is not None:
-            setattr(instance, model_field, result)
-
-        if output_model is not None:
-            for field, value in result.items():
+        if callable(query):
+            query = query(instance)
+        if not query:
+            query = dict()
+        if extra_query:
+            query.update(extra_query)
+
+        output_result, ai_msg = run_llm(
+            prompt, output_model, chat_model, chat_model_name, max_tokens, query=query, extra_tools=tools, **llm_kwargs
+        )
+        # if a result parser is provided, we use this to parse the results into the instance
+        if result_parser:
+            instance = result_parser(instance, output_result, ai_msg)
+        elif output_result and isinstance(output_result, BaseModel):
+            for field, value in output_result.model_dump().items():
                 setattr(instance, field, value)
+    except tuple(BadRequestErrors) as e:  # we silent bad request error because there is nothing we can do about it
+        logger.warning(str(e))
     except tuple(APIStatusErrors) as e:  # for APIStatusError, we let celery retry it
         raise e
     except Exception as e:  # otherwise we log the error and silently fail
@@ -58,15 +73,18 @@ class LLMConfig(Generic[T]):
     def __init__(
         self,
         key: str,
-        prompt: Callable | str,
+        prompt: Callable | str | list[BaseMessage],
         field: str | None = None,
         output_model: "type[BaseModel] | None" = None,
         on_save: bool = True,
         on_condition: Callable | bool | None = None,
-        chat_model: type[BaseChatModel] = ChatOpenAI,
-        chat_model_name: str = "gpt-4o-mini",
+        chat_model: Callable | tuple[type[BaseChatModel], str] = (ChatOpenAI, "gpt-4o-mini"),
         max_tokens: int = 16000,
-        **kwargs,
+        tools: Callable | list[dict[str, str]] | None = None,
+        query: Callable | dict[str, Any] | None = None,
+        llm_kwargs_callback: Callable | None = None,
+        result_parser: Callable | None = None,
+        **llm_kwargs,
     ):
         self.key = key
         self.on_save = on_save
@@ -75,9 +93,12 @@ class LLMConfig(Generic[T]):
         self.prompt = prompt
         self.output_model = output_model
         self.chat_model = chat_model
-        self.chat_model_name = chat_model_name
         self.max_tokens = max_tokens
-        self.kwargs = kwargs
+        self.tools = tools
+        self.query = query
+        self.llm_kwargs_callback = llm_kwargs_callback
+        self.llm_kwargs = llm_kwargs
+        self.result_parser = result_parser
 
     def check_condition(self, instance: T) -> bool:
         if self.on_condition is None:
@@ -86,27 +107,55 @@ class LLMConfig(Generic[T]):
             return self.on_condition(instance)
         return bool(self.on_condition)
 
-    def get_prompt(self, instance: T):
-        prompt = construct_prompt(instance, self.prompt)
+    def _get_prompt(self, instance: T):
+        prompt = self.prompt
+        extra_query = dict()
+        if callable(prompt):
+            prompt = prompt(instance)
         for _, response in add_llm_prompt.send(sender=instance.__class__, instance=instance, key=self.key):
-            prompt.extend(response)
+            remote_prompts, remote_query = response
+            prompt.extend(remote_prompts)
+            extra_query.update(remote_query)
+
+        return prompt, extra_query
+
+    def _get_chat_model(self, instance: T) -> tuple[type[BaseChatModel], str]:
+        if callable(self.chat_model):
+            return self.chat_model(instance)
+        return self.chat_model
+
+    def _get_tools(self, instance: T) -> list[dict[str, Any]]:
+        if callable(self.tools):
+            return self.tools(instance)
+        return self.tools
 
-        return prompt
+    def _get_llm_kwargs(self, instance: T) -> dict[str, Any]:
+        llm_kwargs = self.llm_kwargs
+        if callable(self.llm_kwargs_callback):
+            llm_kwargs = self.llm_kwargs_callback(instance)
+        return llm_kwargs
 
     def schedule(self, instance: T, initial: bool = True):
-        prompt = self.get_prompt(instance)
+        prompt, extra_query = self._get_prompt(instance)
         args = []
         if initial:
             args.append(instance)
+        chat_model, chat_model_name = self._get_chat_model(instance)
+        tools = self._get_tools(instance)
+        llm_kwargs = self._get_llm_kwargs(instance)
         args.extend(
             [
                 prompt,
-                self.chat_model,
-                self.chat_model_name,
+                chat_model,
+                chat_model_name,
                 self.max_tokens,
                 self.field,
                 self.output_model,
-                self.kwargs,
+                tools,
+                self.query,
+                extra_query,
+                llm_kwargs,
+                self.result_parser,
             ]
         )
         return invoke_as_task.s(*args)

wbcore/contrib/ai/llm/mixins.py

@@ -1,10 +1,8 @@
-from contextlib import suppress
-
 from celery import chain, shared_task
 from django.db.models import Model
 
 
-@shared_task
+@shared_task(queue="llm")
 def save_instance_as_task(instance):
     instance.save(_with_llm=False)
 
@@ -25,11 +23,10 @@ class LLMMixin(Model):
                 tasks.append(config.schedule(self, initial=index == 0))
             if tasks:
                 res = chain(*tasks, save_instance_as_task.s())
-                with suppress(Exception):  # we suppress any possible celery exception or pytest issue
-                    if _llm_synchronous:
-                        res.apply()
-                    else:
-                        res.apply_async()
+                if _llm_synchronous:
+                    res.apply()
+                else:
+                    res.apply_async()
 
     class Meta:
         abstract = True

wbcore/contrib/ai/llm/utils.py

@@ -1,45 +1,69 @@
-from typing import Any, Callable
+from contextlib import suppress
+from typing import Any
 
 from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import SystemMessage
 from langchain_core.messages.base import BaseMessage
-from langchain_core.output_parsers import StrOutputParser
+from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
 from langchain_openai import ChatOpenAI
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 
 type Prompt = str | list[BaseMessage] | BaseMessage
 
 
-def construct_prompt(instance: Any, prompt: Prompt | Callable[[Any], Prompt]) -> Prompt:
+def _convert_prompt_to_chat_prompt_template(prompt: Prompt, format_instructions=None) -> ChatPromptTemplate:
+    if isinstance(prompt, BaseMessage):
+        prompt = [prompt]
     if isinstance(prompt, str):
-        return prompt.format(instance)
-
-    elif callable(prompt):
-        return prompt(instance)
-
-    return prompt
+        prompt = SystemMessage(content=prompt)
+    messages = []
+    for index, msg in enumerate(prompt):
+        content = msg.content
+        if index == 0 and format_instructions:
+            if "{format_instructions}" not in content:
+                content += "\n{format_instructions}"
+        messages.append((msg.type, content))
+    return ChatPromptTemplate.from_messages(messages).partial(format_instructions=format_instructions)
 
 
 def run_llm(
-    prompt: Prompt | Callable[[Any], Prompt],
-    instance: Any | None = None,
+    prompt: Prompt,
     output_model: "type[BaseModel] | None" = None,
     chat_model: type[BaseChatModel] = ChatOpenAI,
     chat_model_name: str = "gpt-4o-mini",
     max_tokens: int = 16000,
-    **kwargs,
-) -> str | dict[str, Any]:
-    model = chat_model(model=chat_model_name, max_tokens=max_tokens, **kwargs)  # type: ignore
-    prompt = construct_prompt(instance, prompt)
-    if output_model is None:
-        result = model.invoke(prompt, **kwargs)  # type: ignore
-        parser = StrOutputParser()
-        return parser.invoke(result, **kwargs)
+    extra_tools: list[dict[str, str]] | None = None,
+    query: dict[str, Any] | None = None,
+    **llm_kwargs,
+) -> tuple[BaseModel | None, BaseMessage]:
+    llm = chat_model(model=chat_model_name, max_tokens=max_tokens, **llm_kwargs)  # type: ignore
+    tools = []
+    if not query:
+        query = {}
 
-    structured_model = model.with_structured_output(output_model)
-    result = structured_model.invoke(prompt)  # type: ignore
+    if extra_tools:
+        tools.extend(extra_tools)
+    if output_model:
+        tools.append(output_model)
+    if tools:
+        llm = llm.bind_tools(tools)
+
+    # Set up a parser
+    if output_model:
+        parser = PydanticOutputParser(pydantic_object=output_model)
+    else:
+        parser = StrOutputParser()
+    try:
+        format_instructions = parser.get_format_instructions()
+    except NotImplementedError:
+        format_instructions = None
+    chat_prompt_template = _convert_prompt_to_chat_prompt_template(prompt, format_instructions=format_instructions)
 
-    output = dict()
-    for field, value in result:
-        output[field] = value
+    chain = chat_prompt_template | llm
 
-    return output
+    response = chain.invoke(query)  # type: ignore
+    if output_model:
+        with suppress(ValidationError, IndexError):
+            return output_model.model_validate(response.tool_calls[0]["args"]), response
+    return None, response

wbcore/contrib/authentication/admin.py

@@ -202,11 +202,11 @@ class UserAdmin(admin.ModelAdmin):
             form = self.change_password_form(user)
 
         fieldsets = [(None, {"fields": list(form.base_fields)})]
-        adminForm = admin.helpers.AdminForm(form, fieldsets, {})
+        admin_form = admin.helpers.AdminForm(form, fieldsets, {})
 
         context = {
             "title": _("Change password: %s") % escape(user.get_username()),
-            "adminForm": adminForm,
+            "adminForm": admin_form,
             "form_url": form_url,
             "form": form,
             "is_popup": (IS_POPUP_VAR in request.POST or IS_POPUP_VAR in request.GET),

wbcore/contrib/authentication/factories/__init__.py

@@ -1,3 +1,10 @@
 from .tokens import TokenFactory
-from .users import AuthenticatedPersonFactory, GroupFactory, SuperUserFactory, UserFactory, InternalUserFactory
+from .users import (
+    AuthenticatedPersonFactory,
+    GroupFactory,
+    PermissionFactory,
+    SuperUserFactory,
+    UserFactory,
+    InternalUserFactory,
+)
 from .users_activities import UserActivityFactory

wbcore/contrib/authentication/factories/users.py

@@ -1,5 +1,6 @@
 import factory
 from django.contrib.auth.models import Permission
+from django.contrib.contenttypes.models import ContentType
 from dynamic_preferences.registries import global_preferences_registry
 
 from wbcore.contrib.directory.factories import CompanyFactory, PersonFactory
@@ -90,3 +91,21 @@ class GroupFactory(factory.django.DjangoModelFactory):
 
     class Meta:
         model = Group
+
+
+class ContentTypeFactory(factory.django.DjangoModelFactory):
+    class Meta:
+        model = ContentType
+
+    app_label = "app"
+    model = "model"
+
+
+class PermissionFactory(factory.django.DjangoModelFactory):
+    class Meta:
+        model = Permission
+        django_get_or_create = ("content_type", "codename")
+
+    name = factory.Sequence(lambda n: f"Permission {n}")
+    codename = factory.Sequence(lambda n: f"codename_{n}")
+    content_type = factory.SubFactory(ContentTypeFactory)

wbcore/contrib/authentication/filters.py

@@ -6,11 +6,10 @@ from .models import UserActivity
 
 class UserActivityChartFilter(wb_filters.FilterSet):
     date = wb_filters.DateTimeRangeFilter(
-        method=wb_filters.DateRangeFilter.base_date_range_filter_method,
         label="Date Range",
         required=True,
         clearable=False,
-        default=current_month_date_range,
+        initial=current_month_date_range,
     )
 
     class Meta: