not-again-ai 0.18.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

not_again_ai/data/__init__.py

@@ -1,7 +0,0 @@
-import importlib.util
-
-if importlib.util.find_spec("playwright") is None:
-    raise ImportError(
-        "not_again_ai.data requires the 'data' extra to be installed. "
-        "You can install it using 'pip install not_again_ai[data]'."
-    )

not_again_ai/data/brave_search_api.py

@@ -0,0 +1,203 @@
+import os
+
+import httpx
+from loguru import logger
+from pydantic import BaseModel
+
+
+class SearchWebResult(BaseModel):
+    title: str
+    url: str
+    description: str
+    netloc: str | None = None
+
+
+class SearchWebResults(BaseModel):
+    results: list[SearchWebResult]
+
+
+async def search(
+    query: str,
+    count: int = 20,
+    offset: int = 0,
+    country: str = "US",
+    search_lang: str = "en",
+    ui_lang: str = "en-US",
+    freshness: str | None = None,
+    timezone: str = "America/New_York",
+    state: str = "MA",
+    user_agent: str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36 Edg/134.0.0.",
+) -> SearchWebResults:
+    """
+    Search using Brave Search API.
+
+    Args:
+        query: The search query string
+        count: Number of search results to return (1-20, default 10)
+        offset: Number of search results to skip (default 0)
+        country: Country code for search results (default "US")
+        search_lang: Language for search (default "en")
+        ui_lang: User interface language (default "en-US")
+        freshness: Freshness of results ("pd", "pw", "pm", "py" or YYYY-MM-DDtoYYYY-MM-DD or None)
+        timezone: Timezone for search results (default "America/New_York")
+        state: State for search results (default "MA")
+        user_agent: User agent string for the request (default is a common browser UA)
+
+    Returns:
+        SearchWebResults: A model containing the search results
+
+    Raises:
+        httpx.HTTPError: If the request fails
+        ValueError: If BRAVE_SEARCH_API_KEY is not set
+    """
+    api_key = os.getenv("BRAVE_SEARCH_API_KEY")
+    if not api_key:
+        raise ValueError("BRAVE_SEARCH_API_KEY environment variable is not set")
+
+    url = "https://api.search.brave.com/res/v1/web/search"
+
+    headers = {
+        "Accept": "application/json",
+        "Accept-Encoding": "gzip",
+        "X-Subscription-Token": api_key,
+        "X-Loc-Country": country,
+        "X-Loc-Timezone": timezone,
+        "X-Loc-State": state,
+        "User-Agent": user_agent,
+    }
+
+    params: dict[str, str | int | bool] = {
+        "q": query,
+        "count": count,
+        "offset": offset,
+        "country": country,
+        "search_lang": search_lang,
+        "ui_lang": ui_lang,
+        "text_decorations": False,
+        "spellcheck": False,
+        "units": "imperial",
+        "extra_snippets": False,
+        "safesearch": "off",
+    }
+
+    # Add optional parameters if provided
+    if freshness:
+        params["freshness"] = freshness
+
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, headers=headers, params=params)
+            response.raise_for_status()
+            data = response.json()
+            results_list: list[SearchWebResult] = []
+            for item in data.get("web", {}).get("results", []):
+                result = SearchWebResult(
+                    title=item.get("title", ""),
+                    url=item.get("url", ""),
+                    description=item.get("snippet", ""),
+                    netloc=item.get("meta_url", {}).get("netloc", None),
+                )
+                results_list.append(result)
+            return SearchWebResults(results=results_list)
+
+    except httpx.HTTPError as e:
+        logger.error(f"HTTP error during Brave search: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error during Brave search: {e}")
+        raise
+
+
+class SearchNewsResult(BaseModel):
+    title: str
+    url: str
+    description: str
+    age: str
+    netloc: str | None = None
+
+
+class SearchNewsResults(BaseModel):
+    results: list[SearchNewsResult]
+
+
+async def search_news(
+    query: str,
+    count: int = 20,
+    offset: int = 0,
+    country: str = "US",
+    search_lang: str = "en",
+    ui_lang: str = "en-US",
+    freshness: str | None = None,
+    user_agent: str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36 Edg/134.0.0.",
+) -> SearchNewsResults:
+    """
+    Search news using Brave News Search API.
+
+    Args:
+        query: The search query string
+        count: Number of news results to return (1-20, default 20)
+        offset: Number of search results to skip (default 0)
+        country: Country code for search results (default "US")
+        search_lang: Language for search (default "en")
+        ui_lang: User interface language (default "en-US")
+        freshness: Freshness of results ("pd", "pw", "pm", "py" or YYYY-MM-DDtoYYYY-MM-DD or None)
+        user_agent: User agent string for the request (default is a common browser UA)
+
+    Returns:
+        SearchNewsResults: A model containing the news search results
+
+    Raises:
+        httpx.HTTPError: If the request fails
+        ValueError: If BRAVE_SEARCH_API_KEY is not set
+    """
+    api_key = os.getenv("BRAVE_SEARCH_API_KEY")
+    if not api_key:
+        raise ValueError("BRAVE_SEARCH_API_KEY environment variable is not set")
+
+    url = "https://api.search.brave.com/res/v1/news/search"
+
+    headers = {
+        "Accept": "application/json",
+        "Accept-Encoding": "gzip",
+        "X-Subscription-Token": api_key,
+        "User-Agent": user_agent,
+    }
+
+    params: dict[str, str | int | bool] = {
+        "q": query,
+        "count": count,
+        "offset": offset,
+        "country": country,
+        "search_lang": search_lang,
+        "ui_lang": ui_lang,
+        "spellcheck": False,
+        "safesearch": "off",
+    }
+
+    # Add optional parameters if provided
+    if freshness:
+        params["freshness"] = freshness
+
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url, headers=headers, params=params)
+            response.raise_for_status()
+            data = response.json()
+            results_list: list[SearchNewsResult] = []
+            for item in data.get("results", []):
+                result = SearchNewsResult(
+                    title=item.get("title", ""),
+                    url=item.get("url", ""),
+                    description=item.get("description", ""),
+                    age=item.get("age"),
+                    netloc=item.get("meta_url", {}).get("netloc", None),
+                )
+                results_list.append(result)
+            return SearchNewsResults(results=results_list)
+
+    except httpx.HTTPError as e:
+        logger.error(f"HTTP error during Brave news search: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error during Brave news search: {e}")
+        raise

not_again_ai/data/web.py

@@ -1,56 +1,160 @@
-from loguru import logger
-from playwright.sync_api import Browser, Playwright, sync_playwright
+import asyncio
+import io
+import mimetypes
+from pathlib import Path
+import re
+from urllib.parse import urlparse
 
+from crawl4ai import AsyncWebCrawler, CacheMode
+from crawl4ai.async_configs import BrowserConfig, CrawlerRunConfig
+from crawl4ai.markdown_generation_strategy import DefaultMarkdownGenerator
+import httpx
+from markitdown import MarkItDown, StreamInfo
+from pydantic import BaseModel
 
-def create_browser(headless: bool = True) -> tuple[Playwright, Browser]:
-    """Creates and returns a new Playwright instance and browser.
 
-    Args:
-        headless (bool, optional): Whether to run the browser in headless mode. Defaults to True.
+class Link(BaseModel):
+    url: str
+    text: str
 
-    Returns:
-        tuple[Playwright, Browser]: A tuple containing the Playwright instance and browser.
+
+class URLResult(BaseModel):
+    url: str
+    markdown: str
+    links: list[Link] = []
+
+
+async def _markitdown_bytes_to_str(file_bytes: bytes, filename_extension: str) -> str:
     """
-    pwright = sync_playwright().start()
-    browser = pwright.chromium.launch(
-        headless=headless,
-        chromium_sandbox=False,
-        timeout=15000,
-    )
-    return pwright, browser
+    Convert a file using MarkItDown defaults.
+    """
+    with io.BytesIO(file_bytes) as temp:
+        result = await asyncio.to_thread(
+            MarkItDown(enable_plugins=False).convert,
+            source=temp,
+            stream_info=StreamInfo(extension=filename_extension),
+        )
+        text = result.text_content
+    return text
 
 
-def get_raw_web_content(url: str, browser: Browser | None = None, headless: bool = True) -> str:
-    """Fetches raw web content from a given URL using Playwright.
+def _detect_pdf_extension(url: str) -> bool:
+    """
+    Detect if the URL is a PDF based on its extension.
+    """
+    parsed_url = urlparse(url)
+    filename = Path(parsed_url.path).name
+    return mimetypes.guess_type(filename)[0] == "application/pdf"
 
-    Args:
-        url (str): The URL to fetch content from.
-        browser (Browser | None, optional): An existing browser instance to use. Defaults to None.
-        headless (bool, optional): Whether to run the browser in headless mode. Defaults to True.
 
-    Returns:
-        str: The raw web content.
+def _detect_google_sheets(url: str) -> bool:
+    """
+    Detect if the URL is a Google Sheets document.
     """
-    p = None
-    try:
-        if browser is None:
-            p, browser = create_browser(headless)
+    is_google_sheets = url.startswith("https://docs.google.com/spreadsheets/")
+    return is_google_sheets
+
+
+async def _handle_pdf_content(url: str) -> URLResult:
+    md = MarkItDown(enable_plugins=False)
+    result = md.convert(url)
+    url_result = URLResult(
+        url=url,
+        markdown=result.markdown or "",
+        links=[],
+    )
+    return url_result
+
+
+async def _handle_google_sheets_content(url: str) -> URLResult:
+    """
+    Handle Google Sheets by using the export URL to get the raw content.
+    """
+    edit_pattern = r"https://docs\.google\.com/spreadsheets/d/([a-zA-Z0-9-_]+)/edit"
+    export_pattern = r"https://docs\.google\.com/spreadsheets/d/([a-zA-Z0-9-_]+)/export\?format=csv"
+
+    # Check if it's already an export URL
+    export_match = re.search(export_pattern, url)
+    if export_match:
+        export_url = url
+    else:
+        # Check if it's an edit URL and extract document ID
+        edit_match = re.search(edit_pattern, url)
+        if edit_match:
+            doc_id = edit_match.group(1)
+            export_url = f"https://docs.google.com/spreadsheets/d/{doc_id}/export?format=csv&gid=0"
+        else:
+            return await _handle_web_content(url)
 
-        page = browser.new_page(
-            accept_downloads=False,
-            java_script_enabled=True,
-            viewport={"width": 1366, "height": 768},
-            user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.3",
+    async with httpx.AsyncClient(follow_redirects=True) as client:
+        response = await client.get(export_url)
+        response.raise_for_status()
+        csv_bytes = response.content
+
+    # Convert CSV to markdown using MarkItDown
+    markdown_content = await _markitdown_bytes_to_str(csv_bytes, ".csv")
+
+    url_result = URLResult(
+        url=url,
+        markdown=markdown_content,
+        links=[],
+    )
+    return url_result
+
+
+async def _handle_web_content(url: str) -> URLResult:
+    browser_config = BrowserConfig(
+        browser_type="chromium",
+        headless=True,
+        verbose=False,
+        user_agent_mode="random",
+        java_script_enabled=True,
+    )
+    run_config = CrawlerRunConfig(
+        scan_full_page=True,
+        user_agent_mode="random",
+        cache_mode=CacheMode.DISABLED,
+        markdown_generator=DefaultMarkdownGenerator(),
+    )
+
+    async with AsyncWebCrawler(config=browser_config) as crawler:
+        result = await crawler.arun(
+            url=url,
+            config=run_config,
         )
-        page.goto(url)
-        content = page.content()
-        page.close()
-        return content
-    except Exception as e:
-        logger.error(f"Failed to get web content: {e}")
-        return ""
-    finally:
-        if browser:
-            browser.close()
-        if p:
-            p.stop()
+
+        if result.response_headers.get("content-type") == "application/pdf":
+            return await _handle_pdf_content(url)
+
+    links: list[Link] = []
+    seen_urls: set[str] = set()
+    combined_link_data = result.links.get("internal", []) + result.links.get("external", [])
+    for link_data in combined_link_data:
+        href = link_data.get("href", "")
+        if href and href not in seen_urls:
+            seen_urls.add(href)
+            link = Link(
+                url=href,
+                text=link_data.get("title", "") or link_data.get("text", ""),
+            )
+            links.append(link)
+
+    url_result = URLResult(
+        url=url,
+        markdown=result.markdown or "",
+        links=links,
+    )
+    return url_result
+
+
+async def process_url(url: str) -> URLResult:
+    """
+    Process a URL to extract content and convert it to Markdown and links
+    """
+    if _detect_pdf_extension(url):
+        url_result = await _handle_pdf_content(url)
+    elif _detect_google_sheets(url):
+        url_result = await _handle_google_sheets_content(url)
+    else:
+        url_result = await _handle_web_content(url)
+    return url_result

not_again_ai/llm/chat_completion/interface.py

@@ -2,6 +2,7 @@ from collections.abc import AsyncGenerator, Callable
 from typing import Any
 
 from not_again_ai.llm.chat_completion.providers.anthropic_api import anthropic_chat_completion
+from not_again_ai.llm.chat_completion.providers.gemini_api import gemini_chat_completion
 from not_again_ai.llm.chat_completion.providers.ollama_api import ollama_chat_completion, ollama_chat_completion_stream
 from not_again_ai.llm.chat_completion.providers.openai_api import openai_chat_completion, openai_chat_completion_stream
 from not_again_ai.llm.chat_completion.types import ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse
@@ -16,6 +17,8 @@ def chat_completion(
     - `openai` - OpenAI
     - `azure_openai` - Azure OpenAI
     - `ollama` - Ollama
+    - `anthropic` - Anthropic
+    - `gemini` - Gemini
 
     Args:
         request: Request parameter object
@@ -31,6 +34,8 @@ def chat_completion(
         return ollama_chat_completion(request, client)
     elif provider == "anthropic":
         return anthropic_chat_completion(request, client)
+    elif provider == "gemini":
+        return gemini_chat_completion(request, client)
     else:
         raise ValueError(f"Provider {provider} not supported")
 
@@ -43,8 +48,6 @@ async def chat_completion_stream(
     """Stream a chat completion response from the given provider. Currently supported providers:
     - `openai` - OpenAI
     - `azure_openai` - Azure OpenAI
-    - `ollama` - Ollama
-    - `anthropic` - Anthropic
 
     Args:
         request: Request parameter object

not_again_ai/llm/chat_completion/providers/anthropic_api.py

@@ -103,12 +103,12 @@ def anthropic_chat_completion(request: ChatCompletionRequest, client: Callable[.
         elif tool_choice_value in ["auto", "any"]:
             tool_choice["type"] = "auto"
             if kwargs.get("parallel_tool_calls") is not None:
-                tool_choice["disable_parallel_tool_use"] = str(not kwargs["parallel_tool_calls"])
+                tool_choice["disable_parallel_tool_use"] = not kwargs["parallel_tool_calls"]  # type: ignore
         else:
             tool_choice["name"] = tool_choice_value
             tool_choice["type"] = "tool"
             if kwargs.get("parallel_tool_calls") is not None:
-                tool_choice["disable_parallel_tool_use"] = str(not kwargs["parallel_tool_calls"])
+                tool_choice["disable_parallel_tool_use"] = not kwargs["parallel_tool_calls"]  # type: ignore
         kwargs["tool_choice"] = tool_choice
     kwargs.pop("parallel_tool_calls", None)
 

not_again_ai/llm/chat_completion/providers/gemini_api.py

@@ -0,0 +1,237 @@
+import base64
+from collections.abc import Callable
+import os
+import time
+from typing import Any
+
+from google import genai
+from google.genai import types
+from google.genai.types import FunctionCall, FunctionCallingConfigMode, GenerateContentResponse
+
+from not_again_ai.llm.chat_completion.types import (
+    AssistantMessage,
+    ChatCompletionChoice,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    Function,
+    ImageContent,
+    Role,
+    TextContent,
+    ToolCall,
+)
+
+# This should be all of the options we want to support in types.GenerateContentConfig, that are not handled otherwise
+GEMINI_PARAMETER_MAP = {
+    "max_completion_tokens": "max_output_tokens",
+    "temperature": "temperature",
+    "top_p": "top_p",
+    "top_k": "top_k",
+}
+
+GEMINI_FINISH_REASON_MAP = {
+    "STOP": "stop",
+    "MAX_TOKENS": "max_tokens",
+    "SAFETY": "safety",
+    "RECITATION": "recitation",
+    "LANGUAGE": "language",
+    "OTHER": "other",
+    "BLOCKLIST": "blocklist",
+    "PROHIBITED_CONTENT": "prohibited_content",
+    "SPII": "spii",
+    "MALFORMED_FUNCTION_CALL": "malformed_function_call",
+    "IMAGE_SAFETY": "image_safety",
+}
+
+
+def gemini_chat_completion(request: ChatCompletionRequest, client: Callable[..., Any]) -> ChatCompletionResponse:
+    """Experimental Gemini chat completion function."""
+    # Handle messages
+    # Any system messages need to be removed from messages and concatenated into a single string (in order)
+    system = ""
+    contents = []
+    for message in request.messages:
+        if message.role == "system":
+            # Handle both string content and structured content
+            if isinstance(message.content, str):
+                system += message.content + "\n"
+            else:
+                # If it's a list of content parts, extract text content
+                for part in message.content:
+                    if hasattr(part, "text"):
+                        system += part.text + "\n"
+        elif message.role == "tool":
+            tool_name = message.name if message.name is not None else ""
+            function_response_part = types.Part.from_function_response(
+                name=tool_name,
+                response={"result": message.content},
+            )
+            contents.append(
+                types.Content(
+                    role="user",
+                    parts=[function_response_part],
+                )
+            )
+        elif message.role == "assistant":
+            if message.content and isinstance(message.content, str):
+                contents.append(types.Content(role="model", parts=[types.Part(text=message.content)]))
+            function_parts = []
+            if isinstance(message, AssistantMessage) and message.tool_calls:
+                for tool_call in message.tool_calls:
+                    function_call_part = types.Part(
+                        function_call=FunctionCall(
+                            id=tool_call.id,
+                            name=tool_call.function.name,
+                            args=tool_call.function.arguments,
+                        )
+                    )
+                    function_parts.append(function_call_part)
+            if function_parts:
+                contents.append(types.Content(role="model", parts=function_parts))
+        elif message.role == "user":
+            if isinstance(message.content, str):
+                contents.append(types.Content(role="user", parts=[types.Part(text=message.content)]))
+            elif isinstance(message.content, list):
+                parts = []
+                for part in message.content:
+                    if isinstance(part, TextContent):
+                        parts.append(types.Part(text=part.text))
+                    elif isinstance(part, ImageContent):
+                        # Extract MIME type and data from data URI
+                        uri_parts = part.image_url.url.split(",", 1)
+                        if len(uri_parts) == 2:
+                            mime_type = uri_parts[0].split(":")[1].split(";")[0]
+                            base64_data = uri_parts[1]
+                            image_data = base64.b64decode(base64_data)
+                            parts.append(types.Part.from_bytes(mime_type=mime_type, data=image_data))
+                contents.append(types.Content(role="user", parts=parts))
+
+    kwargs: dict[str, Any] = {}
+    kwargs["contents"] = contents
+    kwargs["model"] = request.model
+    config: dict[str, Any] = {}
+    config["system_instruction"] = system.rstrip()
+    config["automatic_function_calling"] = {"disable": True}
+
+    # Handle the possible tool choice options
+    if request.tool_choice:
+        tool_choice = request.tool_choice
+        if tool_choice == "auto":
+            config["tool_config"] = types.FunctionCallingConfig(mode=FunctionCallingConfigMode.AUTO)
+        elif tool_choice == "any":
+            config["tool_config"] = types.FunctionCallingConfig(mode=FunctionCallingConfigMode.ANY)
+        elif tool_choice == "none":
+            config["tool_config"] = types.FunctionCallingConfig(mode=FunctionCallingConfigMode.NONE)
+        elif isinstance(tool_choice, list):
+            config["tool_config"] = types.FunctionCallingConfig(
+                mode=FunctionCallingConfigMode.ANY, allowed_function_names=tool_choice
+            )
+        elif tool_choice not in (None, "auto", "any", "none"):
+            config["tool_config"] = types.FunctionCallingConfig(
+                mode=FunctionCallingConfigMode.ANY, allowed_function_names=[tool_choice]
+            )
+
+    # Handle tools
+    tools = []
+    for tool in request.tools or []:
+        tools.append(types.Tool(function_declarations=[tool]))  # type: ignore
+    if tools:
+        config["tools"] = tools
+
+    # Everything else defined in GEMINI_PARAMETER_MAP goes into kwargs["config"]
+    request_kwargs = request.model_dump(mode="json", exclude_none=True)
+    for key, value in GEMINI_PARAMETER_MAP.items():
+        if value is not None and key in request_kwargs:
+            config[value] = request_kwargs.pop(key)
+
+    kwargs["config"] = types.GenerateContentConfig(**config)
+
+    start_time = time.time()
+    response: GenerateContentResponse = client(**kwargs)
+    end_time = time.time()
+    response_duration = round(end_time - start_time, 4)
+
+    finish_reason = "other"
+    if response.candidates and response.candidates[0].finish_reason:
+        finish_reason_str = str(response.candidates[0].finish_reason)
+        finish_reason = GEMINI_FINISH_REASON_MAP.get(finish_reason_str, "other")
+
+    tool_calls: list[ToolCall] = []
+    tool_call_objs = response.function_calls
+    if tool_call_objs:
+        for tool_call_obj in tool_call_objs:
+            tool_call_id = tool_call_obj.id if tool_call_obj.id else ""
+            tool_calls.append(
+                ToolCall(
+                    id=tool_call_id,
+                    function=Function(
+                        name=tool_call_obj.name if tool_call_obj.name is not None else "",
+                        arguments=tool_call_obj.args if tool_call_obj.args is not None else {},
+                    ),
+                )
+            )
+
+    assistant_message = ""
+    if (
+        response.candidates
+        and response.candidates[0].content
+        and response.candidates[0].content.parts
+        and response.candidates[0].content.parts[0].text
+    ):
+        assistant_message = response.candidates[0].content.parts[0].text
+
+    choice = ChatCompletionChoice(
+        message=AssistantMessage(
+            role=Role.ASSISTANT,
+            content=assistant_message,
+            tool_calls=tool_calls,
+        ),
+        finish_reason=finish_reason,
+    )
+
+    completion_tokens = 0
+    # Add null check for usage_metadata
+    if response.usage_metadata is not None:
+        if response.usage_metadata.thoughts_token_count:
+            completion_tokens = response.usage_metadata.thoughts_token_count
+        if response.usage_metadata.candidates_token_count:
+            completion_tokens += response.usage_metadata.candidates_token_count
+
+    # Set safe default for prompt_tokens
+    prompt_tokens = 0
+    if response.usage_metadata is not None and response.usage_metadata.prompt_token_count:
+        prompt_tokens = response.usage_metadata.prompt_token_count
+
+    chat_completion_response = ChatCompletionResponse(
+        choices=[choice],
+        completion_tokens=completion_tokens,
+        prompt_tokens=prompt_tokens,
+        response_duration=response_duration,
+    )
+    return chat_completion_response
+
+
+def create_client_callable(client_class: type[genai.Client], **client_args: Any) -> Callable[..., Any]:
+    """Creates a callable that instantiates and uses a Google genai client.
+
+    Args:
+        client_class: The Google genai client class to instantiate
+        **client_args: Arguments to pass to the client constructor
+
+    Returns:
+        A callable that creates a client and returns completion results
+    """
+    filtered_args = {k: v for k, v in client_args.items() if v is not None}
+
+    def client_callable(**kwargs: Any) -> Any:
+        client = client_class(**filtered_args)
+        completion = client.models.generate_content(**kwargs)
+        return completion
+
+    return client_callable
+
+
+def gemini_client(api_key: str | None = None) -> Callable[..., Any]:
+    if not api_key:
+        api_key = os.environ.get("GEMINI_API_KEY")
+    client_callable = create_client_callable(genai.Client, api_key=api_key)
+    return client_callable

not_again_ai/llm/chat_completion/types.py

@@ -138,10 +138,7 @@ class ChatCompletionRequest(BaseModel):
 
 class ChatCompletionChoice(BaseModel):
     message: AssistantMessage
-    finish_reason: Literal[
-        "stop", "length", "tool_calls", "content_filter", "end_turn", "max_tokens", "stop_sequence", "tool_use"
-    ]
-
+    finish_reason: str
     json_message: dict[str, Any] | None = Field(default=None)
     logprobs: list[dict[str, Any] | list[dict[str, Any]]] | None = Field(default=None)
 

not_again_ai/llm/image_gen/__init__.py

@@ -0,0 +1,4 @@
+from not_again_ai.llm.image_gen.interface import create_image
+from not_again_ai.llm.image_gen.types import ImageGenRequest
+
+__all__ = ["ImageGenRequest", "create_image"]

not_again_ai/llm/image_gen/interface.py

@@ -0,0 +1,24 @@
+from collections.abc import Callable
+from typing import Any
+
+from not_again_ai.llm.image_gen.providers.openai_api import openai_create_image
+from not_again_ai.llm.image_gen.types import ImageGenRequest, ImageGenResponse
+
+
+def create_image(request: ImageGenRequest, provider: str, client: Callable[..., Any]) -> ImageGenResponse:
+    """Get a image response from the given provider. Currently supported providers:
+    - `openai` - OpenAI
+    - `azure_openai` - Azure OpenAI
+
+    Args:
+        request: Request parameter object
+        provider: The supported provider name
+        client: Client information, see the provider's implementation for what can be provided
+
+    Returns:
+        ImageGenResponse: The image generation response.
+    """
+    if provider == "openai" or provider == "azure_openai":
+        return openai_create_image(request, client)
+    else:
+        raise ValueError(f"Provider {provider} not supported")

not_again_ai/llm/image_gen/providers/openai_api.py

@@ -0,0 +1,144 @@
+import base64
+from collections.abc import Callable
+import time
+from typing import Any, Literal
+
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from openai import AzureOpenAI, OpenAI
+from openai.types.images_response import ImagesResponse
+
+from not_again_ai.llm.image_gen.types import ImageGenRequest, ImageGenResponse
+
+
+def openai_create_image(request: ImageGenRequest, client: Callable[..., Any]) -> ImageGenResponse:
+    """Create an image using OpenAI API.
+
+    Args:
+        request (ImageGenRequest): The request object containing parameters for image generation.
+        client (Callable[..., Any]): The OpenAI client callable.
+
+    Returns:
+        ImageGenResponse: The response object containing the generated image and metadata.
+    """
+    kwargs = request.model_dump(exclude_none=True)
+    if kwargs.get("images"):
+        kwargs["image"] = kwargs.pop("images", None)
+
+    start_time = time.time()
+    response: ImagesResponse = client(**kwargs)
+    end_time = time.time()
+    response_duration = round(end_time - start_time, 4)
+
+    images: list[bytes] = []
+    if response.data:
+        for data in response.data:
+            images.append(base64.b64decode(data.b64_json or ""))
+
+    input_tokens = response.usage.input_tokens if response.usage else -1
+    output_tokens = response.usage.output_tokens if response.usage else -1
+    input_tokens_details = response.usage.input_tokens_details.to_dict() if response.usage else {}
+    image_gen_response = ImageGenResponse(
+        images=images,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        input_tokens_details=input_tokens_details,
+        response_duration=response_duration,
+    )
+    return image_gen_response
+
+
+def create_client_callable(client_class: type[OpenAI | AzureOpenAI], **client_args: Any) -> Callable[..., Any]:
+    """
+    Creates the correct callable depending on the parameters provided.
+    """
+    filtered_args = {k: v for k, v in client_args.items() if v is not None}
+
+    def client_callable(**kwargs: Any) -> Any:
+        client = client_class(**filtered_args)
+        # If mask or image is not none, use client.images.edit instead of client.images.generate
+        if kwargs.get("mask") or kwargs.get("image"):
+            completion = client.images.edit(**kwargs)
+        else:
+            completion = client.images.generate(**kwargs)
+        return completion
+
+    return client_callable
+
+
+class InvalidOAIAPITypeError(Exception):
+    """Raised when an invalid OAIAPIType string is provided."""
+
+
+def openai_client(
+    api_type: Literal["openai", "azure_openai"] = "openai",
+    api_key: str | None = None,
+    organization: str | None = None,
+    aoai_api_version: str = "2024-06-01",
+    azure_endpoint: str | None = None,
+    timeout: float | None = None,
+    max_retries: int | None = None,
+) -> Callable[..., Any]:
+    """Create an OpenAI or Azure OpenAI client instance based on the specified API type and other provided parameters.
+
+    It is preferred to use RBAC authentication for Azure OpenAI. You must be signed in with the Azure CLI and have correct role assigned.
+    See https://techcommunity.microsoft.com/t5/microsoft-developer-community/using-keyless-authentication-with-azure-openai/ba-p/4111521
+
+    Args:
+        api_type (str, optional): Type of the API to be used. Accepted values are 'openai' or 'azure_openai'.
+            Defaults to 'openai'.
+        api_key (str, optional): The API key to authenticate the client. If not provided,
+            OpenAI automatically uses `OPENAI_API_KEY` from the environment.
+            If provided for Azure OpenAI, it will be used for authentication instead of the Azure AD token provider.
+        organization (str, optional): The ID of the organization. If not provided,
+            OpenAI automotically uses `OPENAI_ORG_ID` from the environment.
+        aoai_api_version (str, optional): Only applicable if using Azure OpenAI https://learn.microsoft.com/azure/ai-services/openai/reference#rest-api-versioning
+        azure_endpoint (str, optional): The endpoint to use for Azure OpenAI.
+        timeout (float, optional): By default requests time out after 10 minutes.
+        max_retries (int, optional): Certain errors are automatically retried 2 times by default,
+            with a short exponential backoff. Connection errors (for example, due to a network connectivity problem),
+            408 Request Timeout, 409 Conflict, 429 Rate Limit, and >=500 Internal errors are all retried by default.
+
+    Returns:
+        Callable[..., Any]: A callable that creates a client and returns completion results
+
+
+    Raises:
+        InvalidOAIAPITypeError: If an invalid API type string is provided.
+        NotImplementedError: If the specified API type is recognized but not yet supported (e.g., 'azure_openai').
+    """
+    if api_type not in ["openai", "azure_openai"]:
+        raise InvalidOAIAPITypeError(f"Invalid OAIAPIType: {api_type}. Must be 'openai' or 'azure_openai'.")
+
+    if api_type == "openai":
+        return create_client_callable(
+            OpenAI,
+            api_key=api_key,
+            organization=organization,
+            timeout=timeout,
+            max_retries=max_retries,
+        )
+    elif api_type == "azure_openai":
+        if api_key:
+            return create_client_callable(
+                AzureOpenAI,
+                api_version=aoai_api_version,
+                azure_endpoint=azure_endpoint,
+                api_key=api_key,
+                timeout=timeout,
+                max_retries=max_retries,
+            )
+        else:
+            azure_credential = DefaultAzureCredential()
+            ad_token_provider = get_bearer_token_provider(
+                azure_credential, "https://cognitiveservices.azure.com/.default"
+            )
+            return create_client_callable(
+                AzureOpenAI,
+                api_version=aoai_api_version,
+                azure_endpoint=azure_endpoint,
+                azure_ad_token_provider=ad_token_provider,
+                timeout=timeout,
+                max_retries=max_retries,
+            )
+    else:
+        raise NotImplementedError(f"API type '{api_type}' is invalid.")

not_again_ai/llm/image_gen/types.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class ImageGenRequest(BaseModel):
+    prompt: str
+    model: str
+    images: list[Path] | None = Field(default=None)
+    mask: Path | None = Field(default=None)
+    n: int = Field(default=1)
+    quality: str | None = Field(default=None)
+    size: str | None = Field(default=None)
+    background: str | None = Field(default=None)
+    moderation: str | None = Field(default=None)
+
+
+class ImageGenResponse(BaseModel):
+    images: list[bytes]
+    input_tokens: int
+    output_tokens: int
+    response_duration: float
+    input_tokens_details: dict[str, Any] | None = Field(default=None)

{not_again_ai-0.18.0.dist-info → not_again_ai-0.20.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: not-again-ai
-Version: 0.18.0
+Version: 0.20.0
 Summary: Designed to once and for all collect all the little things that come up over and over again in AI projects and put them in one place.
 Project-URL: Homepage, https://github.com/DaveCoDev/not-again-ai
 Project-URL: Documentation, https://davecodev.github.io/not-again-ai/
@@ -19,16 +19,18 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Typing :: Typed
 Requires-Python: >=3.11
-Requires-Dist: loguru>=0.7
-Requires-Dist: pydantic>=2.10
+Requires-Dist: loguru<1.0,>=0.7
+Requires-Dist: pydantic<3.0,>=2.11
 Provides-Extra: data
-Requires-Dist: playwright<2.0,>=1.51; extra == 'data'
-Requires-Dist: pytest-playwright<1.0,>=0.7; extra == 'data'
+Requires-Dist: crawl4ai<1.0,>=0.6; extra == 'data'
+Requires-Dist: httpx<1.0,>=0.28; extra == 'data'
+Requires-Dist: markitdown[pdf]==0.1.2; extra == 'data'
 Provides-Extra: llm
-Requires-Dist: anthropic<1.0,>=0.49; extra == 'llm'
+Requires-Dist: anthropic<1.0,>=0.50; extra == 'llm'
 Requires-Dist: azure-identity<2.0,>=1.21; extra == 'llm'
+Requires-Dist: google-genai<2.0,>1.12; extra == 'llm'
 Requires-Dist: ollama<1.0,>=0.4; extra == 'llm'
-Requires-Dist: openai<2.0,>=1.68; extra == 'llm'
+Requires-Dist: openai<2.0,>=1.76; extra == 'llm'
 Requires-Dist: python-liquid<3.0,>=2.0; extra == 'llm'
 Requires-Dist: tiktoken<1.0,>=0.9; extra == 'llm'
 Provides-Extra: statistics
@@ -62,7 +64,7 @@ It is encouraged to also **a)** use this as a template for your own Python packa
 **b)** instead of installing the package, copy and paste functions into your own projects. 
 We make this easier by limiting the number of dependencies and use an MIT license.
 
-**Documentation** available within individual **[notebooks](notebooks)**, docstrings within the source, or auto-generated at [DaveCoDev.github.io/not-again-ai/](https://DaveCoDev.github.io/not-again-ai/).
+**Documentation** available within individual **[notebooks](notebooks)** or docstrings within the source code.
 
 # Installation
 
@@ -82,7 +84,9 @@ The package is split into subpackages, so you can install only the parts you nee
 
 ### Data
 1. `pip install not_again_ai[data]`
-1. `playwright install` to download the browser binaries.
+1. `crawl4ai-setup` to run crawl4ai post-installation setup.
+1. Set the `BRAVE_SEARCH_API_KEY` environment variable to use the Brave Search API for web data extraction.
+   1. Get the API key from https://api-dashboard.search.brave.com/app/keys. You must have at least the Free "Data for Search" subscription.
 
 
 ### LLM
@@ -138,7 +142,7 @@ all machines that use the project, both during development and in production.
 To install all dependencies into an isolated virtual environment:
 
 ```shell
-uv sync --all-extras
+uv sync --all-extras --all-groups
 ```
 
 To upgrade all dependencies to their latest versions:
@@ -311,3 +315,5 @@ Default settings are configured in [`.vscode/settings.json`](./.vscode/settings.
 
 # Attributions
 [python-blueprint](https://github.com/johnthagen/python-blueprint) for the Python package skeleton.
+
+This project uses Crawl4AI (https://github.com/unclecode/crawl4ai) for web data extraction.

{not_again_ai-0.18.0.dist-info → not_again_ai-0.20.0.dist-info}/RECORD

@@ -3,14 +3,16 @@ not_again_ai/py.typed,sha256=UaCuPFa3H8UAakbt-5G8SPacldTOGvJv18pPjUJ5gDY,93
 not_again_ai/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 not_again_ai/base/file_system.py,sha256=KNQmacO4Q__CQuq2oPzWrg3rQO48n3evglc9bNiP7KM,949
 not_again_ai/base/parallel.py,sha256=fcYhKBYBWvob84iKp3O93wvFFdXeidljZsShgBLTNGA,3448
-not_again_ai/data/__init__.py,sha256=1jF6mwvtB2PT7IEc3xpbRtZm3g3Lyf8zUqH4AEE4qlQ,244
-not_again_ai/data/web.py,sha256=wjx9cc33jcoJBGonYCIpwygPBFOwz7F-dx_ominmbnI,1838
+not_again_ai/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+not_again_ai/data/brave_search_api.py,sha256=cTH7smRf8ITPpcQEcoIEvJo6VAST9_mg7FSL88xdGMc,6603
+not_again_ai/data/web.py,sha256=LLWM5SkgI3-ILhtXYYuDm_eVnij3uS4tL059RyMq0lU,4737
 not_again_ai/llm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 not_again_ai/llm/chat_completion/__init__.py,sha256=HozawvdRkTFgq8XR16GJUHN1ukEa4Ya68wVPVrl-afs,250
-not_again_ai/llm/chat_completion/interface.py,sha256=xRZXQ75dxrkt5WNtOTtrAa2Oy4ZB-PG2WihW9FBmW-s,2525
-not_again_ai/llm/chat_completion/types.py,sha256=Z_pQjVK_7rEvAE2fj5srKxHPFJBLPV8e0iCFibnzT7M,5596
+not_again_ai/llm/chat_completion/interface.py,sha256=OU6ghG7RlveahkHZWdRHFg0uzbSrSh2Dz7u5-4rrypA,2700
+not_again_ai/llm/chat_completion/types.py,sha256=0pBo1Fgm__JU3NyMShGouAIolcANPpTfXn8WJHODlQw,5472
 not_again_ai/llm/chat_completion/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-not_again_ai/llm/chat_completion/providers/anthropic_api.py,sha256=_-NPc5pfhsRwwy-GYc1vAiyc0agmGLyo5_7-mcPEnBU,6189
+not_again_ai/llm/chat_completion/providers/anthropic_api.py,sha256=Eix3_GgQvDyPr6pKfSDrfqRg_bTb0pqMI8fdQrx9e84,6211
+not_again_ai/llm/chat_completion/providers/gemini_api.py,sha256=ovyTssfN3achMr2Laa2Hu557CjYN6o7UkO5IGXg6lzk,9461
 not_again_ai/llm/chat_completion/providers/ollama_api.py,sha256=Puo2eE2VynvZOoqrUlNYtPgRGCRMVa8syc3TfBxS1hs,10661
 not_again_ai/llm/chat_completion/providers/openai_api.py,sha256=1wdeV50KYX_KIf2uofsICKYoHVSvj4kTRpS1Vuw3NSQ,17887
 not_again_ai/llm/embedding/__init__.py,sha256=wscUfROukvw0M0vYccfaVTdXV0P-eICAT5mqM0LaHHc,182
@@ -19,6 +21,11 @@ not_again_ai/llm/embedding/types.py,sha256=J4FFLx35Aow2kOaafDReeY9cUNqhWMjaAk5gX
 not_again_ai/llm/embedding/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 not_again_ai/llm/embedding/providers/ollama_api.py,sha256=m-OCis9WAUT2baGsGVPzejlive40eSNyO6tHmPh6joM,3201
 not_again_ai/llm/embedding/providers/openai_api.py,sha256=JFFqbq0O5snIEnr9VESdp5xehikQBPbs7nwyE6acFsY,5441
+not_again_ai/llm/image_gen/__init__.py,sha256=v31PgYdTxMQRRxXPFl40BW5Y8RSHrZuwabuD-yC9gfI,170
+not_again_ai/llm/image_gen/interface.py,sha256=XGE0aDvQwe-aWRuGNLMECO6KnMiK8qLv2APvr0hZ0tY,930
+not_again_ai/llm/image_gen/types.py,sha256=Qhdov5azWwmmbqqE3Ln7t-Fb_Ipyp8r3z0_80omyASc,672
+not_again_ai/llm/image_gen/providers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+not_again_ai/llm/image_gen/providers/openai_api.py,sha256=3IEhdId1UU_imvDebytTt0dSCjEEyPHf2o4vVt6RqSE,6198
 not_again_ai/llm/prompting/__init__.py,sha256=7YnHro1yH01FLGnao27WyqQDFjNYf9npE5UxoR9YrUU,84
 not_again_ai/llm/prompting/compile_prompt.py,sha256=uBn655yTiQ325z1CUgnkU2k7ICIvaYRJOm50B7w2lSs,4683
 not_again_ai/llm/prompting/interface.py,sha256=SMKYabmu3zTWbEDukU6aLU_JQ88apeBWWOF_qZ0s3ww,1783
@@ -33,7 +40,7 @@ not_again_ai/viz/distributions.py,sha256=OyWwJaNI6lMRm_iSrhq-CORLNvXfeuLSgDtVo3u
 not_again_ai/viz/scatterplot.py,sha256=5CUOWeknbBOaZPeX9oPin5sBkRKEwk8qeFH45R-9LlY,2292
 not_again_ai/viz/time_series.py,sha256=pOGZqXp_2nd6nKo-PUQNCtmMh__69jxQ6bQibTGLwZA,5212
 not_again_ai/viz/utils.py,sha256=hN7gwxtBt3U6jQni2K8j5m5pCXpaJDoNzGhBBikEU28,238
-not_again_ai-0.18.0.dist-info/METADATA,sha256=n41TWZaLvs_XPNbSEQojju9DI4nPhkjE055xX7ZJGjQ,12021
-not_again_ai-0.18.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-not_again_ai-0.18.0.dist-info/licenses/LICENSE,sha256=btjOgNGpp-ux5xOo1Gx1MddxeWtT9sof3s3Nui29QfA,1071
-not_again_ai-0.18.0.dist-info/RECORD,,
+not_again_ai-0.20.0.dist-info/METADATA,sha256=sYV6N2dmYarZnbjIqLoCrjfrsa8m1G-1rYxX3XV89tI,12389
+not_again_ai-0.20.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+not_again_ai-0.20.0.dist-info/licenses/LICENSE,sha256=btjOgNGpp-ux5xOo1Gx1MddxeWtT9sof3s3Nui29QfA,1071
+not_again_ai-0.20.0.dist-info/RECORD,,