feat: summarizer (#20)

Author: JarbasAl

ovos_solver_openai_persona/engines.py

@@ -2,19 +2,27 @@
 from typing import Optional, Iterable, List, Dict
 
 import requests
-from ovos_plugin_manager.templates.solvers import QuestionSolver
-from ovos_utils.log import LOG
 
 from ovos_plugin_manager.templates.solvers import ChatMessageSolver
+from ovos_plugin_manager.templates.solvers import QuestionSolver
+from ovos_plugin_manager.templates.language import LanguageTranslator, LanguageDetector
+from ovos_utils.log import LOG
 
 MessageList = List[Dict[str, str]]  # for typing
 
-class OpenAICompletionsSolver(QuestionSolver):
-    enable_tx = False
-    priority = 25
 
-    def __init__(self, config=None):
-        super().__init__(config)
+class OpenAICompletionsSolver(QuestionSolver):
+    def __init__(self, config=None,
+                 translator: Optional[LanguageTranslator] = None,
+                 detector: Optional[LanguageDetector] = None,
+                 priority: int = 50,
+                 enable_tx: bool = False,
+                 enable_cache: bool = False,
+                 internal_lang: Optional[str] = None):
+        super().__init__(config=config, translator=translator,
+                         detector=detector, priority=priority,
+                         enable_tx=enable_tx, enable_cache=enable_cache,
+                         internal_lang=internal_lang)
         self.api_url = f"{self.config.get('api_url', 'https://api.openai.com/v1')}/completions"
         self.engine = self.config.get("model", "text-davinci-002")  # "ada" cheaper and faster, "davinci" better
         self.stop_token = "<|im_end|>"
@@ -78,11 +86,17 @@ def post_process_sentence(text: str) -> str:
 
 
 class OpenAIChatCompletionsSolver(ChatMessageSolver):
-    enable_tx = False
-    priority = 25
-
-    def __init__(self, config=None):
-        super().__init__(config)
+    def __init__(self, config=None,
+                 translator: Optional[LanguageTranslator] = None,
+                 detector: Optional[LanguageDetector] = None,
+                 priority: int = 25,
+                 enable_tx: bool = False,
+                 enable_cache: bool = False,
+                 internal_lang: Optional[str] = None):
+        super().__init__(config=config, translator=translator,
+                         detector=detector, priority=priority,
+                         enable_tx=enable_tx, enable_cache=enable_cache,
+                         internal_lang=internal_lang)
         self.api_url = f"{self.config.get('api_url', 'https://api.openai.com/v1')}/chat/completions"
         self.engine = self.config.get("model", "gpt-4o-mini")  # "ada" cheaper and faster, "davinci" better
         self.stop_token = "<|im_end|>"
@@ -91,7 +105,7 @@ def __init__(self, config=None):
             LOG.error("key not set in config")
             raise ValueError("key must be set")
         self.memory = config.get("enable_memory", True)
-        self.max_utts = config.get("memory_size", 15)
+        self.max_utts = config.get("memory_size", 5)
         self.qa_pairs = []  # tuple of q+a
         self.initial_prompt = config.get("initial_prompt", "You are a helpful assistant.")
 
@@ -267,4 +281,3 @@ def get_spoken_answer(self, query: str,
         messages = self.get_messages(query)
         # just for api compat since it's a subclass, shouldn't be directly used
         return self.continue_chat(messages=messages, lang=lang, units=units)
-

ovos_solver_openai_persona/summarizer.py

@@ -0,0 +1,41 @@
+from ovos_solver_openai_persona.engines import OpenAIChatCompletionsSolver
+from ovos_plugin_manager.templates.solvers import TldrSolver
+from ovos_plugin_manager.templates.language import LanguageTranslator, LanguageDetector
+from typing import Dict, Optional
+
+
+class OpenAISummarizer(TldrSolver):
+    TEMPLATE = """Your task is to summarize the text into a suitable format.
+Answer in plaintext with no formatting, 2 paragraphs long at most. 
+Focus on the most important information.
+---------------------
+{content}
+"""
+    def __init__(self, config: Optional[Dict] = None,
+                 translator: Optional[LanguageTranslator] = None,
+                 detector: Optional[LanguageDetector] = None,
+                 priority: int = 50,
+                 enable_tx: bool = False,
+                 enable_cache: bool = False,
+                 internal_lang: Optional[str] = None):
+        super().__init__(config=config, translator=translator,
+                         detector=detector, priority=priority,
+                         enable_tx=enable_tx, enable_cache=enable_cache,
+                         internal_lang=internal_lang)
+        self.llm = OpenAIChatCompletionsSolver(config=config, translator=translator,
+                                               detector=detector, priority=priority,
+                                               enable_tx=enable_tx, enable_cache=enable_cache,
+                                               internal_lang=internal_lang)
+        self.prompt_template = self.config.get("prompt_template") or self.TEMPLATE
+
+
+    def get_tldr(self, document: str, lang: Optional[str] = None) -> str:
+        """
+        Summarize the provided document.
+
+        :param document: The text of the document to summarize, assured to be in the default language.
+        :param lang: Optional language code.
+        :return: A summary of the provided document.
+        """
+        prompt = self.prompt_template.format(content=document)
+        return self.llm.get_spoken_answer(prompt, lang)

setup.py

@@ -50,6 +50,7 @@ def get_version():
 PERSONA_ENTRY_POINT = 'Remote Llama=ovos_solver_openai_persona:LLAMA_DEMO'
 PLUGIN_ENTRY_POINT = 'ovos-solver-openai-persona-plugin=ovos_solver_openai_persona:OpenAIPersonaSolver'
 DIALOG_PLUGIN_ENTRY_POINT = 'ovos-dialog-transformer-openai-plugin=ovos_solver_openai_persona.dialog_transformers:OpenAIDialogTransformer'
+SUMMARIZER_ENTRY_POINT = 'ovos-summarizer-openai-plugin=ovos_solver_openai_persona.summarizer:OpenAISummarizer'
 
 
 setup(
@@ -66,6 +67,7 @@ def get_version():
     entry_points={
         'neon.plugin.solver': PLUGIN_ENTRY_POINT,
         "opm.transformer.dialog": DIALOG_PLUGIN_ENTRY_POINT,
+        'opm.solver.summarization': SUMMARIZER_ENTRY_POINT,
         "opm.plugin.persona": PERSONA_ENTRY_POINT
     },
     install_requires=required("requirements.txt"),