pygpt-net 2.1.37__py3-none-any.whl → 2.4.48__py3-none-any.whl

pygpt_net/core/idx/chat.py

@@ -6,18 +6,25 @@
 # GitHub:  https://github.com/szczyglis-dev/py-gpt   #
 # MIT License                                        #
 # Created By  : Marcin Szczygliński                  #
-# Updated Date: 2024.03.13 01:00:00                  #
+# Updated Date: 2025.01.16 01:00:00                  #
 # ================================================== #
 
 import json
+from typing import Optional, Dict, Any, List
 
 from llama_index.core.llms import ChatMessage, MessageRole
 from llama_index.core.prompts import ChatPromptTemplate
 from llama_index.core.memory import ChatMemoryBuffer
 
+from pygpt_net.core.types import (
+    MODE_CHAT,
+    MODE_AGENT_LLAMA,
+)
+from pygpt_net.core.bridge.context import BridgeContext
+from pygpt_net.item.model import ModelItem
 from pygpt_net.item.ctx import CtxItem
+
 from .context import Context
-from pygpt_net.item.model import ModelItem
 
 
 class Chat:
@@ -31,49 +38,83 @@ class Chat:
         self.storage = storage
         self.context = Context(window)
 
-    def call(self, **kwargs) -> bool:
+    def call(
+            self,
+            context: BridgeContext,
+            extra: Optional[Dict[str, Any]] = None
+    ) -> bool:
         """
-        Call chat or query mode
+        Call chat, complete or query mode
 
-        :param kwargs: keyword arguments
+        :param context: Bridge context
+        :param extra: Extra arguments
         :return: True if success
         """
-        model = kwargs.get("model", None)
-        idx_raw = kwargs.get("idx_raw", False)  # raw mode
+        model = context.model
+        idx_mode = context.idx_mode  # mode
+
         if model is None or not isinstance(model, ModelItem):  # check if model is provided
             raise Exception("Model config not provided")
 
-        if idx_raw:  # query index (raw mode)
-            return self.raw_query(**kwargs)
+        if idx_mode == "query":  # query index only (raw mode)
+            return self.raw_query(
+                context=context,
+                extra=extra,
+            )
+        elif idx_mode == "retrieval":  # retrieval mode
+            return self.retrieval(
+                context=context,
+                extra=extra,
+            )
 
         # if not raw, check if chat mode is available
-        if "chat" in model.llama_index['mode']:
-            return self.chat(**kwargs)
+        if MODE_CHAT in model.llama_index['mode']:
+            return self.chat(
+                context=context,
+                extra=extra,
+            )
         else:
-            return self.query(**kwargs)  # if not, use query mode
-
-    def raw_query(self, **kwargs) -> bool:
+            return self.query(
+                context=context,
+                extra=extra,
+            )  # if not, use query mode
+
+    def raw_query(
+            self,
+            context: BridgeContext,
+            extra: Optional[Dict[str, Any]] = None
+    ) -> bool:
         """
         Raw query mode
 
-        :param kwargs: keyword arguments
+        :param context: Bridge context
+        :param extra: Extra arguments
         :return: True if success
         """
-        return self.query(**kwargs)
+        return self.query(
+            context=context,
+            extra=extra,
+        )
 
-    def query(self, **kwargs) -> bool:
+    def query(
+            self,
+            context: BridgeContext,
+            extra: Optional[Dict[str, Any]] = None
+    ) -> bool:
         """
         Query index mode (no chat, only single query) and append results to context
 
-        :param kwargs: keyword arguments
+        :param context: Bridge context
+        :param extra: Extra arguments
         :return: True if success
         """
-        ctx = kwargs.get("ctx", CtxItem())
-        idx = kwargs.get("idx", "base")
-        model = kwargs.get("model", None)
-        system_prompt = kwargs.get("system_prompt_raw", None)  # raw system prompt, without plugin ads
-        stream = kwargs.get("stream", False)
-        query = ctx.input
+        idx = context.idx
+        model = context.model
+        system_prompt = context.system_prompt_raw  # get raw system prompt, without plugin addons
+        stream = context.stream
+        ctx = context.ctx
+        query = ctx.input  # user input
+        verbose = self.window.core.config.get("log.llama", False)
 
         if model is None or not isinstance(model, ModelItem):
             raise Exception("Model config not provided")
@@ -85,12 +126,7 @@ class Chat:
             model.id,
         ))
 
-        # check if index exists
-        if not self.storage.exists(idx):
-            raise Exception("Index not prepared")
-
-        context = self.window.core.idx.llm.get_service_context(model=model)
-        index = self.storage.get(idx, service_context=context)  # get index
+        index, llm = self.get_index(idx, model)
         input_tokens = self.window.core.tokens.from_llama_messages(
             query,
             [],
@@ -101,22 +137,26 @@ class Chat:
         if tpl is not None:
             self.log("Query index with custom prompt: {}...".format(system_prompt))
             response = index.as_query_engine(
+                llm=llm,
                 streaming=stream,
                 text_qa_template=tpl,
+                verbose=verbose,
             ).query(query)  # query with custom sys prompt
         else:
             response = index.as_query_engine(
+                llm=llm,
                 streaming=stream,
+                verbose=verbose,
             ).query(query)  # query with default prompt
 
         if response:
             if stream:
-                ctx.add_doc_meta(response.metadata)  # store metadata
+                ctx.add_doc_meta(self.get_metadata(response.source_nodes))  # store metadata
                 ctx.stream = response.response_gen
                 ctx.input_tokens = input_tokens
                 ctx.set_output("", "")
             else:
-                ctx.add_doc_meta(response.metadata)  # store metadata
+                ctx.add_doc_meta(self.get_metadata(response.source_nodes))  # store metadata
                 ctx.input_tokens = input_tokens
                 ctx.output_tokens = self.window.core.tokens.from_llama_messages(
                     response.response,
@@ -124,79 +164,210 @@ class Chat:
                     model.id,
                 )  # calc from response
                 ctx.set_output(str(response.response), "")
-            self.log("Found in metadata: {}".format(response.metadata))
             return True
         return False
 
-    def chat(self, **kwargs) -> bool:
+    def retrieval(
+            self,
+            context: BridgeContext,
+            extra: Optional[Dict[str, Any]] = None
+    ) -> bool:
         """
-        Chat mode (conversation, using context from index) and append result to the context
+        Retrieve documents from index only
 
-        :param kwargs: keyword arguments
+        :param context: Bridge context
+        :param extra: Extra arguments
         :return: True if success
         """
-        ctx = kwargs.get("ctx", CtxItem())
-        idx = kwargs.get("idx", "base")
-        model = kwargs.get("model", None)
-        system_prompt = kwargs.get("system_prompt", None)
-        stream = kwargs.get("stream", False)
-        query = ctx.input
+        idx = context.idx
+        model = context.model
+        ctx = context.ctx
+        query = ctx.input  # user input
+        verbose = self.window.core.config.get("log.llama", False)
+
+        self.log("Retrieval...")
+        self.log("Idx: {}, retrieve only: {}".format(
+            idx,
+            query,
+        ))
+
+        index, llm = self.get_index(idx, model)
+        retriever = index.as_retriever()
+        nodes = retriever.retrieve(query)
+        outputs = []
+        self.log("Retrieved {} nodes...".format(len(nodes)))
+        for node in nodes:
+            outputs.append({
+                "text": node.text,
+                "score": node.score,
+            })
+        if outputs:
+            response = ""
+            for output in outputs:
+                response += "**Score: {}**\n\n{}".format(output["score"], output["text"])
+                if output != outputs[-1]:
+                    response += "\n\n-------\n\n"
+            ctx.set_output(response)
+            ctx.add_doc_meta(self.get_metadata(nodes))
+        return True
+
+    def chat(
+            self,
+            context: BridgeContext,
+            extra: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Chat mode (conversation, using context from index) and append result to the context
+
+        :param context: Bridge context
+        :param extra: Extra arguments
+        """
+        idx = context.idx
+        model = context.model
+        system_prompt = context.system_prompt  # get final system prompt
+        stream = context.stream
+        ctx = context.ctx
+        query = ctx.input  # user input
+        chat_mode = self.window.core.config.get("llama.idx.chat.mode")
+        use_index = True
+        verbose = self.window.core.config.get("log.llama", False)
+
+        if idx is None or idx == "_":
+            chat_mode = "simple"  # do not use query engine if no index
+            use_index = False
 
         if model is None or not isinstance(model, ModelItem):
             raise Exception("Model config not provided")
 
         self.log("Chat with index...")
-        self.log("Idx: {}, query: {}, model: {}".format(
+        self.log("Idx: {}, query: {}, chat_mode: {}, model: {}".format(
             idx,
             query,
+            chat_mode,
             model.id,
         ))
 
-        # check if index exists
-        if not self.storage.exists(idx):
-            raise Exception("Index not prepared")
+        # use index only if idx is not empty, otherwise use only LLM
+        index = None
+        if use_index:
+            index, llm = self.get_index(idx, model)
+        else:
+            llm = self.window.core.idx.llm.get(model)
 
-        context = self.window.core.idx.llm.get_service_context(model=model)
-        index = self.storage.get(idx, service_context=context)  # get index
+        # if multimodal support, try to get multimodal provider
+        # if model.is_multimodal():
+            # llm = self.window.core.idx.llm.get(model, multimodal=True)  # get multimodal LLM model
 
         # append context from DB
-        history = self.context.get_messages(ctx.input, system_prompt)
-        memory = self.get_memory_buffer(history, context.llm)
+        history = self.context.get_messages(
+            ctx.input,
+            system_prompt,
+            context.history,
+        )
+        memory = self.get_memory_buffer(history, llm)
         input_tokens = self.window.core.tokens.from_llama_messages(
             query,
             history,
             model.id,
         )
-        chat_engine = index.as_chat_engine(
-            chat_mode="context",
-            memory=memory,
-            system_prompt=system_prompt,
-        )
-        if stream:
-            response = chat_engine.stream_chat(query)
+
+        if use_index:
+            # CMD: commands are applied to system prompt here
+            # index as query engine
+            chat_engine = index.as_chat_engine(
+                llm=llm,
+                chat_mode=chat_mode,
+                memory=memory,
+                verbose=verbose,
+                system_prompt=system_prompt,
+            )
+            if stream:
+                response = chat_engine.stream_chat(query)
+            else:
+                response = chat_engine.chat(query)
         else:
-            response = chat_engine.chat(query)
+            # CMD: commands are applied to system prompt here
+            # prepare tools (native calls if enabled)
+            tools = self.window.core.agents.tools.prepare(context, extra)
 
-        if response:
+            # without index, use LLM as chat engine
+            history.insert(0, self.context.add_system(system_prompt))
+            history.append(self.context.add_user(query))
             if stream:
-                ctx.stream = response.response_gen
-                ctx.input_tokens = input_tokens
-                ctx.set_output("", "")
-                # no metadata here in agent chat stream response!
+                if hasattr(llm, "stream_chat_with_tools"):
+                    response = llm.stream_chat_with_tools(
+                        tools=tools,
+                        messages=history,
+                    )
+                else:
+                    response = llm.stream_chat(
+                        messages=history,
+                    )
             else:
-                ctx.add_doc_meta(response.metadata)  # store metadata
-                ctx.input_tokens = input_tokens
-                ctx.output_tokens = self.window.core.tokens.from_llama_messages(
-                    response.response,
-                    [],
-                    model.id,
-                )  # calc from response
-                ctx.set_output(str(response.response), "")
-                self.log("Found in metadata: {}".format(response.metadata))
+                if hasattr(llm, "chat_with_tools"):
+                    response = llm.chat_with_tools(
+                        tools=tools,
+                        messages=history,
+                    )
+                else:
+                    response = llm.chat(
+                        messages=history,
+                    )
+
+        # handle response
+        if response:
+            if use_index:
+                # from index as query engine
+                if stream:
+                    ctx.stream = response.response_gen
+                    ctx.input_tokens = input_tokens
+                    ctx.set_output("", "")
+                    ctx.add_doc_meta(self.get_metadata(response.source_nodes))  # store metadata
+                else:
+                    ctx.input_tokens = input_tokens
+                    ctx.output_tokens = self.window.core.tokens.from_llama_messages(
+                        response.response,
+                        [],
+                        model.id,
+                    )  # calc from response
+                    output = response.response
+                    if output is None:
+                        output = ""
+                    ctx.set_output(output, "")
+                    ctx.add_doc_meta(self.get_metadata(response.source_nodes))  # store metadata
+            else:
+                # from LLM directly
+                if stream:
+                    # tools handled in stream output controller
+                    ctx.stream = response  # chunk is in response.delta
+                    ctx.input_tokens = input_tokens
+                    ctx.set_output("", "")
+                else:
+                    # unpack tool calls
+                    tool_calls = llm.get_tool_calls_from_response(
+                        response, error_on_no_tool_call=False
+                    )
+                    ctx.tool_calls = self.window.core.command.unpack_tool_calls_from_llama(tool_calls)
+                    output = response.message.content
+                    if output is None:
+                        output = ""
+                    ctx.set_output(output, "")
+                    ctx.input_tokens = input_tokens
+                    ctx.output_tokens = self.window.core.tokens.from_llama_messages(
+                        output,
+                        [],
+                        model.id,
+                    ) # calc from response
             return True
         return False
 
-    def query_file(self, ctx: CtxItem, path: str, query: str, model: ModelItem = None) -> str:
+    def query_file(
+            self,
+            ctx: CtxItem,
+            path: str,
+            query: str,
+            model: Optional[ModelItem] = None
+    ) -> str:
         """
         Query file using temp index (created on the fly)
 
@@ -208,8 +379,9 @@ class Chat:
         """
         if model is None:
             model = self.window.core.models.from_defaults()
-        context = self.window.core.idx.llm.get_service_context(model=model)
-        tmp_id, index = self.storage.get_tmp(path, service_context=context)  # get or create tmp index
+
+        llm, embed_model = self.window.core.idx.llm.get_service_context(model=model)
+        tmp_id, index = self.storage.get_tmp(path, llm, embed_model)  # get or create tmp index
 
         idx = "tmp:{}".format(path)  # tmp index id
         self.log("Indexing to temporary in-memory index: {}...".format(idx))
@@ -227,10 +399,11 @@ class Chat:
         if len(files) > 0:
             self.log("Querying temporary in-memory index: {}...".format(idx))
             response = index.as_query_engine(
+                llm=llm,
                 streaming=False,
             ).query(query)  # query with default prompt
             if response:
-                ctx.add_doc_meta(response.metadata)  # store metadata
+                ctx.add_doc_meta(self.get_metadata(response.source_nodes))  # store metadata
                 output = response.response
 
         # clean tmp index
@@ -239,7 +412,15 @@ class Chat:
         self.log("Returning response: {}".format(output))
         return output
 
-    def query_web(self, ctx: CtxItem, type: str, url: str, args: dict, query: str, model: ModelItem = None) -> str:
+    def query_web(
+            self,
+            ctx: CtxItem,
+            type: str,
+            url: str,
+            args: Dict[str, Any],
+            query: str,
+            model: Optional[ModelItem] = None
+    ) -> str:
         """
         Query web using temp index (created on the fly)
 
@@ -259,8 +440,8 @@ class Chat:
         id = json.dumps(parts)
         if model is None:
             model = self.window.core.models.from_defaults()
-        context = self.window.core.idx.llm.get_service_context(model=model)
-        tmp_id, index = self.storage.get_tmp(id, service_context=context)  # get or create tmp index
+        llm, embed_model = self.window.core.idx.llm.get_service_context(model=model)
+        tmp_id, index = self.storage.get_tmp(id, llm, embed_model)  # get or create tmp index
 
         idx = "tmp:{}".format(id)  # tmp index id
         self.log("Indexing to temporary in-memory index: {}...".format(idx))
@@ -280,19 +461,109 @@ class Chat:
         if num > 0:
             self.log("Querying temporary in-memory index: {}...".format(idx))
             response = index.as_query_engine(
+                llm=llm,
                 streaming=False,
             ).query(query)  # query with default prompt
             if response:
-                ctx.add_doc_meta(response.metadata)  # store metadata
+                ctx.add_doc_meta(self.get_metadata(response.source_nodes))  # store metadata
                 output = response.response
 
-        # clean tmp index# clean tmp index
+        # clean tmp index
         self.log("Removing temporary in-memory index: {} ({})...".format(idx, tmp_id))
         self.storage.clean_tmp(tmp_id)  # clean memory
         self.log("Returning response: {}...".format(output))
         return output
 
-    def get_memory_buffer(self, history: list, llm = None) -> ChatMemoryBuffer:
+    def query_attachment(
+            self,
+            query: str,
+            path: str,
+            model: Optional[ModelItem] = None,
+            history: Optional[List[CtxItem]] = None,
+            verbose: bool = False,
+    ) -> str:
+        """
+        Query attachment
+
+        :param query: query
+        :param path: path to index
+        :param model: model
+        :param history: chat history
+        :param verbose: verbose mode
+        :return: response
+        """
+        if model is None:
+            model = self.window.core.models.from_defaults()
+        llm, embed_model = self.window.core.idx.llm.get_service_context(model=model)
+        index = self.storage.get_ctx_idx(path, llm, embed_model)
+
+        # 1. try to retrieve directly from index
+        retriever = index.as_retriever()
+        nodes = retriever.retrieve(query)
+        response = ""
+        score = 0
+        for node in nodes:
+            if node.score > 0.5:
+                score = node.score
+                response = node.text
+                break
+        output = ""
+        if response:
+            output = str(response)
+            if verbose:
+                print("Found using retrieval, {} (score: {})".format(output, score))
+        else:
+            if verbose:
+                print("Not found using retrieval, trying with query engine...")
+            history = self.context.get_messages(
+                query,
+                "",
+                history,
+            )
+            memory = self.get_memory_buffer(history, llm)
+            response = index.as_chat_engine(
+                llm=llm,
+                streaming=False,
+                memory=memory,
+            ).chat(query)
+            if response:
+                output = str(response.response)
+        return output
+
+    def query_retrieval(
+            self,
+            query: str,
+            idx: str,
+            model: Optional[ModelItem] = None
+    ) -> str:
+        """
+        Query attachment
+
+        :param query: query
+        :param idx: index id
+        :param model: model
+        :return: response
+        """
+        if model is None:
+            model = self.window.core.models.from_defaults()
+        index, llm = self.get_index(idx, model)
+        retriever = index.as_retriever()
+        nodes = retriever.retrieve(query)
+        response = ""
+        for node in nodes:
+            if node.score > 0.5:
+                response = node.text
+                break
+        output = ""
+        if response:
+            output = str(response)
+        return output
+
+    def get_memory_buffer(
+            self,
+            history: List[ChatMessage],
+            llm = None
+    ) -> ChatMemoryBuffer:
         """
         Get memory buffer
 
@@ -305,7 +576,10 @@ class Chat:
             llm=llm,
         )
 
-    def get_custom_prompt(self, prompt: str = None) -> ChatPromptTemplate or None:
+    def get_custom_prompt(
+            self,
+            prompt: Optional[str] = None
+    ) -> Optional[ChatPromptTemplate]:
         """
         Get custom prompt template if sys prompt is not empty
 
@@ -334,12 +608,92 @@ class Chat:
         ]
         return ChatPromptTemplate(qa_msgs)
 
+    def get_index(
+            self,
+            idx: str,
+            model: ModelItem
+    ):
+        """
+        Get index instance
+
+        :param idx: idx name (id)
+        :param model: model instance
+        :return:
+        """
+        # check if index exists
+        if not self.storage.exists(idx):
+            if idx is None:
+                # create empty in memory idx
+                llm, embed_model = self.window.core.idx.llm.get_service_context(model=model)
+                index = self.storage.index_from_empty()
+                return index, llm
+            # raise Exception("Index not prepared")
+
+        llm, embed_model = self.window.core.idx.llm.get_service_context(model=model)
+        index = self.storage.get(idx, llm, embed_model)  # get index
+        return index, llm
+
+    def get_metadata(
+            self,
+            source_nodes: Optional[list]
+    ) -> Dict[str, Any]:
+        """
+        Get metadata from source nodes
+
+        :param source_nodes: source nodes
+        :return: metadata
+        """
+        if (source_nodes is None
+                or not isinstance(source_nodes, list)
+                or len(source_nodes) == 0):
+            return {}
+        metadata = {}
+        i = 1
+        max = 3
+        min_score = 0.3
+        for node in source_nodes:
+            if hasattr(node, "id_"):
+                id = node.id_
+                if node.metadata is not None:
+                    score = node.get_score()
+                    if score > min_score:
+                        metadata[id] = node.metadata
+                        metadata[id]["score"] = score
+                        i += 1
+                        if i > max:
+                            break
+        return metadata
+
+    def parse_metadata(
+            self,
+            metadata: Optional[Dict]
+    ) -> Dict[str, Any]:
+        """
+        Parse metadata
+
+        :param metadata: metadata
+        :return: metadata
+        """
+        if (metadata is None
+                or not isinstance(metadata, dict)
+                or len(metadata) == 0):
+            return {}
+        meta = {}
+        for node in metadata:
+            doc_id = list(node.keys())[0]
+            meta[doc_id] = node[doc_id]
+            break
+        return meta
+
     def log(self, msg: str):
         """
         Log info message
 
         :param msg: message
         """
+        # disabled logging for thread safety
+        if self.window.core.config.get("mode") == MODE_AGENT_LLAMA:
+            return
         is_log = False
         if self.window.core.config.has("log.llama") \
                 and self.window.core.config.get("log.llama"):
@@ -347,3 +701,4 @@ class Chat:
         self.window.core.debug.info(msg, not is_log)
         if is_log:
             print("[LLAMA-INDEX] {}".format(msg))
+        self.window.idx_logger_message.emit(msg)