refactor + add text generation webui internal client

README.md

@@ -44,3 +44,4 @@ TODO
 - prompt - add prompt enhancer
 - prompt - add prompt compression using LLMLingua
 - llm provider - optionally use NuExtract text to json model to get structured response(like instead of embedding the instruction to the base llm, take the resp and send it to nuextract)
+- llm provider - text generation web ui - support multiple instances hosted in diff locations

src/ezpyai/_constants.py

@@ -1,15 +1,25 @@
-from typing import Dict
+LIB_NAME: str = "ezpyai"
 
-_LIB_NAME: str = "ezpyai"
+# ENVIRONMENT VARIABLES
+ENV_VAR_NAME_OPENAI_API_KEY: str = "OPENAI_API_KEY"
+ENV_VAR_NAME_OPENAI_ORGANIZATION: str = "OPENAI_ORGANIZATION"
+ENV_VAR_NAME_OPENAI_PROJECT: str = "OPENAI_PROJECT"
+ENV_VAR_NAME_TEXT_GENERATION_WEBUI_BASE_URL: str = "TEXT_GENERATION_WEBUI_BASE_URL"
+ENV_VAR_NAME_TEXT_GENERATION_WEBUI_API_KEY: str = "TEXT_GENERATION_WEBUI_API_KEY"
 
-_ENV_VAR_NAME_OPENAI_API_KEY: str = "OPENAI_API_KEY"
-_ENV_VAR_NAME_OPENAI_ORGANIZATION: str = "OPENAI_ORGANIZATION"
-_ENV_VAR_NAME_OPENAI_PROJECT: str = "OPENAI_PROJECT"
-_ENV_VAR_NAME_TEXT_GENERATION_WEBUI_API_KEY: str = "TEXT_GENERATION_WEBUI_API_KEY"
-_ENV_VAR_NAME_TEXT_GENERATION_WEBUI_BASE_URL: str = "TEXT_GENERATION_WEBUI_BASE_URL"
+# DICTIONARY KEYS
+DICT_KEY_ID: str = "id"
+DICT_KEY_METADATA: str = "metadata"
+DICT_KEY_CONTENT: str = "content"
+DICT_KEY_SUMMARY: str = "summary"
 
+# HTTP METHODS
+HTTP_METHOD_GET: str = "GET"
+HTTP_METHOD_POST: str = "POST"
 
-_DICT_KEY_ID: str = "id"
-_DICT_KEY_METADATA: str = "metadata"
-_DICT_KEY_CONTENT: str = "content"
-_DICT_KEY_SUMMARY: str = "summary"
+# HTTP CONTENT TYPES
+HTTP_CONTENT_TYPE_JSON: str = "application/json"
+
+# HTTP HEADERS
+HTTP_HEADER_CONTENT_TYPE: str = "Content-Type"
+HTTP_HEADER_AUTHORIZATION: str = "Authorization"

src/ezpyai/_logger.py

@@ -1,4 +1,4 @@
 import logging
-from ezpyai._constants import _LIB_NAME
+from ezpyai._constants import LIB_NAME
 
-logger = logging.getLogger(_LIB_NAME)
+logger = logging.getLogger(LIB_NAME)

src/ezpyai/llm/knowledge/_knowledge_gatherer.py

@@ -14,7 +14,7 @@
 from PyPDF2 import PdfReader
 from docx import Document
 from ezpyai._logger import logger
-from ezpyai._constants import _DICT_KEY_SUMMARY
+from ezpyai._constants import DICT_KEY_SUMMARY
 from ezpyai.llm.providers._llm_provider import LLMProvider
 from ezpyai.llm.prompt import Prompt, get_summarizer_prompt
 from ezpyai.llm.knowledge.knowledge_item import KnowledgeItem
@@ -45,7 +45,12 @@ class KnowledgeGatherer:
     """
 
     def __init__(self, summarizer: LLMProvider = None) -> None:
-        """Initialize the KnowledgeGatherer with an empty _items dictionary."""
+        """
+        Initialize the KnowledgeGatherer with an empty _items dictionary.
+
+        Args:
+            summarizer (LLMProvider, optional): The LLMProvider hosting the summarizer model to use for knowledge collection. Defaults to None.
+        """
 
         self._items: Dict[str, KnowledgeItem] = {}
         self._summarizer: LLMProvider = summarizer
@@ -114,8 +119,8 @@ def _summarize(self, knowledge_item: KnowledgeItem) -> None:
         prompt: Prompt = get_summarizer_prompt(knowledge_item.content)
 
         knowledge_item.summary = self._summarizer.get_structured_response(
-            prompt, response_format={_DICT_KEY_SUMMARY: ""}
-        )[_DICT_KEY_SUMMARY]
+            prompt, response_format={DICT_KEY_SUMMARY: ""}
+        )[DICT_KEY_SUMMARY]
 
         logger.debug(f"Summarized knowledge item: {knowledge_item}")
 

src/ezpyai/llm/knowledge/chroma_db.py

@@ -3,7 +3,7 @@
 
 from typing import Dict, List
 from ezpyai._logger import logger
-from ezpyai._constants import _DICT_KEY_SUMMARY
+from ezpyai._constants import DICT_KEY_SUMMARY
 from ezpyai.llm.providers._llm_provider import LLMProvider
 from ezpyai.llm.knowledge._knowledge_db import BaseKnowledgeDB
 from ezpyai.llm.knowledge._knowledge_gatherer import KnowledgeGatherer
@@ -88,7 +88,7 @@ def store(
             logger.debug(f"Pre-processing item: {knowledge_item}")
 
             metadata = knowledge_item.metadata
-            metadata[_DICT_KEY_SUMMARY] = knowledge_item.summary
+            metadata[DICT_KEY_SUMMARY] = knowledge_item.summary
 
             document_ids.append(knowledge_item.id)
             documents.append(knowledge_item.content)
@@ -141,8 +141,8 @@ def search(
 
         for i in range(len(documents)):
             summary = ""
-            if _DICT_KEY_SUMMARY in metadatas[i]:
-                summary = metadatas[i].pop(_DICT_KEY_SUMMARY, "")
+            if DICT_KEY_SUMMARY in metadatas[i]:
+                summary = metadatas[i].pop(DICT_KEY_SUMMARY, "")
 
             knowledge_items.append(
                 KnowledgeItem(

src/ezpyai/llm/knowledge/knowledge_item.py

@@ -1,9 +1,9 @@
 from typing import Dict
 from ezpyai._constants import (
-    _DICT_KEY_ID,
-    _DICT_KEY_METADATA,
-    _DICT_KEY_CONTENT,
-    _DICT_KEY_SUMMARY,
+    DICT_KEY_ID,
+    DICT_KEY_METADATA,
+    DICT_KEY_CONTENT,
+    DICT_KEY_SUMMARY,
 )
 
 
@@ -38,8 +38,8 @@ def __str__(self):
 
     def to_dict(self) -> Dict[str, str]:
         return {
-            _DICT_KEY_ID: self.id,
-            _DICT_KEY_METADATA: self.metadata,
-            _DICT_KEY_CONTENT: self.content,
-            _DICT_KEY_SUMMARY: self.summary,
+            DICT_KEY_ID: self.id,
+            DICT_KEY_METADATA: self.metadata,
+            DICT_KEY_CONTENT: self.content,
+            DICT_KEY_SUMMARY: self.summary,
         }