Merge pull request #38 from imranq2/add-test-for-prompt-helper

Add test for prompt helper. Remove logging unless LOG_INPUT_AND_OUTPUT is set

Dockerfile

@@ -124,6 +124,7 @@ USER appuser
 # The number of workers can be controlled using the NUM_WORKERS environment variable
 # Otherwise the number of workers for uvicorn (using the multiprocessing worker) is  chosen based on these guidelines:
 # (https://sentry.io/answers/number-of-uvicorn-workers-needed-in-production/)
+# basically (cores * threads + 1)
 CMD ["sh", "-c", "\
     # Get CPU info \
     CORE_COUNT=$(nproc) && \

docker-compose-openwebui.yml

@@ -25,13 +25,13 @@ services:
 #      TITLE_GENERATION_PROMPT_TEMPLATE: "What do you want to know about your health?"
       ENABLE_IMAGE_GENERATION: True
       IMAGE_GENERATION_ENGINE: "openai"
-      IMAGES_OPENAI_API_BASE_URL: "http://dev:5000/api/v1"
+      IMAGES_OPENAI_API_BASE_URL: "http://localhost:5050/api/v1"
       IMAGES_OPENAI_API_KEY: "bedrock"
       IMAGE_GENERATION_MODEL: "amazon.titan-image-generator-v2:0"
       LOAD_PROMPT_FROM_LANGCHAIN_HUB: False
       WEBUI_FAVICON_URL: "http://dev:5000/favicon.png"
       WEBUI_NAME: "b.ai"
-      WEBUI_BANNERS: '[{"id": "1", "type":"info", "title":"Welcome to b.ai", "content":"Welcome to b.well AI (https://www.cnn.com)", "dismissible": false, "timestamp": 100000}]'
+      WEBUI_BANNERS: '[{"id": "1", "type":"info", "title":"Welcome to b.well AI", "content":"Welcome to b.well AI (https://icanbwell.atlassian.net/wiki/spaces/BRG/pages/5092442117/AI+Resource+Guide)", "dismissible": false, "timestamp": 100000},{"id": "2", "type":"info", "title":"Help", "content":"Type \"help\" if you need help", "dismissible": false, "timestamp": 100000}]'
   #    volumes:
 #      - ./openwebui/backend/data:/app/backend/data
 #      - ./openwebui-config/config.json:/app/backend/data/config.json:cached

docker-compose.yml

@@ -27,6 +27,7 @@ services:
       NUM_WORKERS: 1
       AWS_REGION: 'us-east-1'
       HELP_KEYWORDS: "help;/help;aid"
+      LOG_INPUT_AND_OUTPUT: 1
     env_file: docker.env
     ports:
       - '5050:5000'

language_model_gateway/configs/chat_completions/prompt_helper.json

@@ -0,0 +1,50 @@
+{
+  "$schema": "https://raw.githubusercontent.com/imranq2/language_model_gateway/main/language_model_gateway/configs/config_schema.json",
+  "id": "prompt_helper",
+  "name": "Prompt Helper",
+  "description": "This models helps you craft the best prompt. Specify your goal and constraints, and I will help you create a detailed system prompt to guide a language model in completing the task effectively.",
+  "type": "langchain",
+  "model": {
+    "provider": "bedrock",
+    "model": "us.anthropic.claude-3-5-haiku-20241022-v1:0"
+  },
+  "system_prompts": [
+    {
+      "role": "system",
+      "content": "Given a task description or existing prompt, produce a detailed system prompt to guide a language model in completing the task effectively.\n\n# Guidelines\n\n- Understand the Task: Grasp the main objective, goals, requirements, constraints, and expected output.\n- Minimal Changes: If an existing prompt is provided, improve it only if it's simple. For complex prompts, enhance clarity and add missing elements without altering the original structure.\n- Reasoning Before Conclusions**: Encourage reasoning steps before any conclusions are reached. ATTENTION! If the user provides examples where the reasoning happens afterward, REVERSE the order! NEVER START EXAMPLES WITH CONCLUSIONS!\n    - Reasoning Order: Call out reasoning portions of the prompt and conclusion parts (specific fields by name). For each, determine the ORDER in which this is done, and whether it needs to be reversed.\n    - Conclusion, classifications, or results should ALWAYS appear last.\n- Examples: Include high-quality examples if helpful, using placeholders [in brackets] for complex elements.\n   - What kinds of examples may need to be included, how many, and whether they are complex enough to benefit from placeholders.\n- Clarity and Conciseness: Use clear, specific language. Avoid unnecessary instructions or bland statements.\n- Formatting: Use markdown features for readability. DO NOT USE ``` CODE BLOCKS UNLESS SPECIFICALLY REQUESTED.\n- Preserve User Content: If the input task or prompt includes extensive guidelines or examples, preserve them entirely, or as closely as possible. If they are vague, consider breaking down into sub-steps. Keep any details, guidelines, examples, variables, or placeholders provided by the user.\n- Constants: DO include constants in the prompt, as they are not susceptible to prompt injection. Such as guides, rubrics, and examples.\n- Output Format: Explicitly the most appropriate output format, in detail. This should include length and syntax (e.g. short sentence, paragraph, JSON, etc.)\n    - For tasks outputting well-defined or structured data (classification, JSON, etc.) bias toward outputting a JSON.\n    - JSON should never be wrapped in code blocks (```) unless explicitly requested.\n\nThe final prompt you output should adhere to the following structure below. Do not include any additional commentary, only output the completed system prompt. SPECIFICALLY, do not include any additional messages at the start or end of the prompt. (e.g. no \"---\")\n\n[Concise instruction describing the task - this should be the first line in the prompt, no section header]\n\n[Additional details as needed.]\n\n[Optional sections with headings or bullet points for detailed steps.]\n\n# Steps [optional]\n\n[optional: a detailed breakdown of the steps necessary to accomplish the task]\n\n# Output Format\n\n[Specifically call out how the output should be formatted, be it response length, structure e.g. JSON, markdown, etc]\n\n# Examples [optional]\n\n[Optional: 1-3 well-defined examples with placeholders if necessary. Clearly mark where examples start and end, and what the input and output are. User placeholders as necessary.]\n[If the examples are shorter than what a realistic example is expected to be, make a reference with () explaining how real examples should be longer / shorter / different. AND USE PLACEHOLDERS! ]\n\n# Notes [optional]\n\n[optional: edge cases, details, and an area to call or repeat out specific important considerations]"
+    },
+    {
+      "role": "system",
+      "content": "The user will provide a Task, Goal, or Current Prompt"
+    }
+  ],
+  "model_parameters": [
+    {
+      "key": "temperature",
+      "value": 0.5
+    }
+  ],
+  "headers": [
+    {
+      "key": "Authorization",
+      "value": "Bearer OPENAI_API_KEY"
+    }
+  ],
+  "tools": [
+    {
+      "name": "current_date",
+      "parameters": [
+        {
+          "key": "format",
+          "value": "YYYY-MM-DD"
+        }
+      ]
+    }
+  ],
+  "example_prompts": [
+    {
+      "role": "user",
+      "content": "Specify your goal and constraints, and I will help you create a detailed system prompt to guide a language model in completing the task effectively."
+    }
+  ]
+}

language_model_gateway/gateway/image_generation/aws_image_generator.py

@@ -2,6 +2,7 @@
 import base64
 import json
 import logging
+import os
 from concurrent.futures.thread import ThreadPoolExecutor
 from typing import override, Dict, Any
 
@@ -39,7 +40,8 @@ async def generate_image_async(
         self, prompt: str, style: str = "natural", image_size: str = "1024x1024"
     ) -> bytes:
         """Generate an image using Titan Image Generator"""
-        logger.info(f"Generating image for prompt: {prompt}")
+        if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+            logger.info(f"Generating image for prompt: {prompt}")
 
         request_body = {
             "textToImageParams": {"text": prompt},
@@ -72,7 +74,8 @@ async def generate_image_async(
             # Convert base64 to bytes
             image_data = base64.b64decode(base64_image)
 
-            logger.info(f"Image generated successfully for prompt: {prompt}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Image generated successfully for prompt: {prompt}")
             return image_data
 
         except Exception as e:

language_model_gateway/gateway/managers/chat_completion_manager.py

@@ -118,7 +118,8 @@ async def chat_completions(
         if help_response is not None:
             return help_response
 
-        logger.info(f"Running chat completion for {chat_request}")
+        if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+            logger.info(f"Running chat completion for {chat_request}")
         # Use the provider to get the completions
         return await provider.chat_completions(
             model_config=model_config, headers=headers, chat_request=chat_request
@@ -175,9 +176,10 @@ def handle_help_prompt(
             )
 
         last_message_content: str = cast(str, user_messages[-1]["content"])
-        logger.info(
-            f"Last message content: {last_message_content}, type: {type(last_message_content)}"
-        )
+        if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+            logger.info(
+                f"Last message content: {last_message_content}, type: {type(last_message_content)}"
+            )
 
         help_keywords: List[str] = os.environ.get("HELP_KEYWORDS", "help").split(";")
         if (
@@ -220,7 +222,8 @@ def handle_help_prompt(
                 created=int(time.time()),
                 object="chat.completion",
             )
-            logger.info(f"Returning help response: {chat_response.model_dump()}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Returning help response: {chat_response.model_dump()}")
             if chat_request.get("stream"):
 
                 async def foo(

language_model_gateway/gateway/providers/image_generation_provider.py

@@ -62,7 +62,8 @@ async def generate_image_async(
             image_generation_request.get("response_format")
         )
 
-        logger.info(f"image_generation_request: {image_generation_request}")
+        if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+            logger.info(f"image_generation_request: {image_generation_request}")
 
         model: Union[str, ImageModel, None] | NotGiven = image_generation_request.get(
             "model"

language_model_gateway/gateway/providers/openai_chat_completions_provider.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 from os import environ
 from random import randint
 from typing import Any, Dict, AsyncGenerator
@@ -95,7 +96,8 @@ async def chat_completions(
                     content=f"Error validating response: {e}. url: {agent_url}\n{response_text}",
                     status_code=500,
                 )
-            logger.info(f"Non-streaming response {request_id}: {response}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Non-streaming response {request_id}: {response}")
             return JSONResponse(content=response.model_dump())
 
     async def get_streaming_response_async(
@@ -143,12 +145,13 @@ async def _stream_resp_async_generator(
                     data: str = sse.data
                     i += 1
 
-                    if logger.isEnabledFor(logging.DEBUG):
-                        logger.debug(
-                            f"----- Received data from stream {i} {event} {type(data)} ------"
-                        )
-                        logger.debug(data)
-                        logger.debug(
-                            f"----- End data from stream {i} {event} {type(data)} ------"
-                        )
+                    if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                        if logger.isEnabledFor(logging.DEBUG):
+                            logger.debug(
+                                f"----- Received data from stream {i} {event} {type(data)} ------"
+                            )
+                            logger.debug(data)
+                            logger.debug(
+                                f"----- End data from stream {i} {event} {type(data)} ------"
+                            )
                     yield f"data: {data}\n\n"

language_model_gateway/gateway/tools/google_search_tool.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import os
 import random
 from os import environ
 from typing import Optional, Dict, Any, List, cast, Type, Literal, Tuple
@@ -82,9 +83,10 @@ async def _make_request(self, url: str, params: Dict[str, Any]) -> Dict[str, Any
 
         while True:
             try:
-                logger.info(
-                    f"Running Google search with query {params['q']}.  Params: {params}.  Retry count: {retry_count}"
-                )
+                if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                    logger.info(
+                        f"Running Google search with query {params['q']}.  Params: {params}.  Retry count: {retry_count}"
+                    )
 
                 response = await self._client.get(url, params=params)
 
@@ -143,7 +145,8 @@ async def _arun(self, query: str) -> Tuple[str, str]:
                     snippets.append(f"- {result['snippet']} ({result.get('link')})")
 
             response: str = "\n".join(snippets)
-            logger.info(f"Google Search results: {response}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Google Search results: {response}")
             return response, f'GoogleSearchTool: Searched Google for "{query}"'
         except Exception as e:
             logger.exception(e, stack_info=True)

language_model_gateway/gateway/tools/python_repl_tool.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Type
 
 from langchain_core.tools import BaseTool
@@ -23,9 +24,11 @@ async def _arun(self, query: str) -> str:
         """Async implementation of the tool (in this case, just calls _run)"""
         try:
             python_repl = PythonREPL()
-            logger.info(f"Running Python Repl with query: {query}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Running Python Repl with query: {query}")
             result: str = python_repl.run(command=query)
-            logger.info(f"Python Repl result: {result}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Python Repl result: {result}")
             return result
         except Exception as e:
             logger.error(f"Error running Python Repl: {e}")
@@ -35,9 +38,11 @@ async def _arun(self, query: str) -> str:
     def _run(self, query: str) -> str:
         try:
             python_repl = PythonREPL()
-            logger.info(f"Running Python Repl with query: {query}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Running Python Repl with query: {query}")
             result: str = python_repl.run(command=query)
-            logger.info(f"Python Repl result: {result}")
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(f"Python Repl result: {result}")
             return result
         except Exception as e:
             logger.error(f"Error running Python Repl: {e}")

language_model_gateway/gateway/tools/scraping_bee_web_scraper_tool.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Optional, Dict, Type, Tuple, Literal
 
 import httpx
@@ -78,13 +79,17 @@ async def _async_scrape(self, *, url: str, query: Optional[str]) -> Optional[str
 
         try:
             async with httpx.AsyncClient() as client:
-                logger.info(f"Scraping {url} with ScrapingBee with params: {params}")
+                if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                    logger.info(
+                        f"Scraping {url} with ScrapingBee with params: {params}"
+                    )
                 response = await client.get(self.base_url, params=params, timeout=30.0)
 
                 if response.status_code == 200:
-                    logger.info(
-                        f"====== Scraped {url} ======\n{response.text}\n====== End of Scraped Content ======"
-                    )
+                    if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                        logger.info(
+                            f"====== Scraped {url} ======\n{response.text}\n====== End of Scraped Content ======"
+                        )
                     return response.text
                 else:
                     logger.error(

language_model_gateway/gateway/tools/url_to_markdown_tool.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Type, Literal, Tuple
 
 import httpx
@@ -54,9 +55,10 @@ async def _arun(self, url: str) -> Tuple[str, str]:
             content: str = await HtmlToMarkdownConverter.get_markdown_from_html_async(
                 html_content=html_content
             )
-            logger.info(
-                f"====== Scraped {url} ======\n{content}\n====== End of Scraped Markdown ======"
-            )
+            if os.environ.get("LOG_INPUT_AND_OUTPUT", "0") == "1":
+                logger.info(
+                    f"====== Scraped {url} ======\n{content}\n====== End of Scraped Markdown ======"
+                )
             return content, f"URLToMarkdownTool: Scraped content from <{url}> "
         except Exception as e:
             return (

tests/gateway/test_chat_anthropic_image_generation.py

@@ -3,39 +3,25 @@
 from openai import AsyncOpenAI
 from openai.types import ImagesResponse
 
-from language_model_gateway.container.simple_container import SimpleContainer
-from language_model_gateway.gateway.api_container import get_container_async
-from language_model_gateway.gateway.image_generation.image_generator_factory import (
-    ImageGeneratorFactory,
-)
-from language_model_gateway.gateway.models.model_factory import ModelFactory
-from language_model_gateway.gateway.utilities.environment_reader import (
-    EnvironmentReader,
-)
-from tests.gateway.mocks.mock_chat_model import MockChatModel
-from tests.gateway.mocks.mock_image_generator import MockImageGenerator
-from tests.gateway.mocks.mock_image_generator_factory import MockImageGeneratorFactory
-from tests.gateway.mocks.mock_model_factory import MockModelFactory
-
 
 @pytest.mark.asyncio
 async def test_chat_anthropic_image_generation(async_client: httpx.AsyncClient) -> None:
     print("")
 
-    if not EnvironmentReader.is_environment_variable_set("RUN_TESTS_WITH_REAL_LLM"):
-        test_container: SimpleContainer = await get_container_async()
-        test_container.register(
-            ModelFactory,
-            lambda c: MockModelFactory(
-                fn_get_model=lambda chat_model_config: MockChatModel(
-                    fn_get_response=lambda messages: "His first name is Barack"
-                )
-            ),
-        )
-        test_container.register(
-            ImageGeneratorFactory,
-            lambda c: MockImageGeneratorFactory(image_generator=MockImageGenerator()),
-        )
+    # if not EnvironmentReader.is_environment_variable_set("RUN_TESTS_WITH_REAL_LLM"):
+    #     test_container: SimpleContainer = await get_container_async()
+    #     test_container.register(
+    #         ModelFactory,
+    #         lambda c: MockModelFactory(
+    #             fn_get_model=lambda chat_model_config: MockChatModel(
+    #                 fn_get_response=lambda messages: "His first name is Barack"
+    #             )
+    #         ),
+    #     )
+    #     test_container.register(
+    #         ImageGeneratorFactory,
+    #         lambda c: MockImageGeneratorFactory(image_generator=MockImageGenerator()),
+    #     )
 
     # init client and connect to localhost server
     client = AsyncOpenAI(