feat: remove _generate for llm

api/app.py

@@ -14,7 +14,7 @@ from flask import Flask, request, Response, session
 import flask_login
 from flask_cors import CORS
 
-from extensions import ext_session, ext_celery, ext_sentry, ext_redis, ext_login, ext_vector_store, ext_migrate, \
+from extensions import ext_session, ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
     ext_database, ext_storage
 from extensions.ext_database import db
 from extensions.ext_login import login_manager
@@ -79,7 +79,6 @@ def initialize_extensions(app):
     ext_database.init_app(app)
     ext_migrate.init(app, db)
     ext_redis.init_app(app)
-    ext_vector_store.init_app(app)
     ext_storage.init_app(app)
     ext_celery.init_app(app)
     ext_session.init_app(app)

api/core/callback_handler/llm_callback_handler.py

@@ -53,37 +53,12 @@ class LLMCallbackHandler(BaseCallbackHandler):
     ) -> None:
         self.start_at = time.perf_counter()
 
-        # todo chat serialized maybe deprecated in future
-        if 'Chat' in serialized['name']:
-            real_prompts = []
-            messages = []
-            for prompt in prompts:
-                role, content = prompt.split(': ', maxsplit=1)
-                if role == 'human':
-                    role = 'user'
-                    message = HumanMessage(content=content)
-                elif role == 'ai':
-                    role = 'assistant'
-                    message = AIMessage(content=content)
-                else:
-                    message = SystemMessage(content=content)
-
-                real_prompt = {
-                    "role": role,
-                    "text": content
-                }
-                real_prompts.append(real_prompt)
-                messages.append(message)
-
-            self.llm_message.prompt = real_prompts
-            self.llm_message.prompt_tokens = self.llm.get_messages_tokens(messages)
-        else:
-            self.llm_message.prompt = [{
-                "role": 'user',
-                "text": prompts[0]
-            }]
+        self.llm_message.prompt = [{
+            "role": 'user',
+            "text": prompts[0]
+        }]
 
-            self.llm_message.prompt_tokens = self.llm.get_num_tokens(prompts[0])
+        self.llm_message.prompt_tokens = self.llm.get_num_tokens(prompts[0])
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         end_at = time.perf_counter()

api/core/callback_handler/std_out_callback_handler.py

@@ -19,6 +19,7 @@ class DifyStdOutCallbackHandler(BaseCallbackHandler):
             messages: List[List[BaseMessage]],
             **kwargs: Any
     ) -> Any:
+        print_text("\n[on_chat_model_start]\n", color='blue')
         for sub_messages in messages:
             for sub_message in sub_messages:
                 print_text(str(sub_message) + "\n", color='blue')
@@ -28,12 +29,7 @@ class DifyStdOutCallbackHandler(BaseCallbackHandler):
     ) -> None:
         """Print out the prompts."""
         print_text("\n[on_llm_start]\n", color='blue')
-
-        if 'Chat' in serialized['name']:
-            for prompt in prompts:
-                print_text(prompt + "\n", color='blue')
-        else:
-            print_text(prompts[0] + "\n", color='blue')
+        print_text(prompts[0] + "\n", color='blue')
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Do nothing."""

api/core/index/keyword_table_index/keyword_table_index.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 from typing import Any, List, Optional, Dict
 
 from langchain.schema import Document, BaseRetriever
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, Extra
 
 from core.index.base import BaseIndex
 from core.index.keyword_table_index.jieba_keyword_table_handler import JiebaKeywordTableHandler
@@ -170,6 +170,12 @@ class KeywordTableRetriever(BaseRetriever, BaseModel):
     index: KeywordTableIndex
     search_kwargs: dict = Field(default_factory=dict)
 
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
     def get_relevant_documents(self, query: str) -> List[Document]:
         """Get documents relevant for a query.
 

api/core/llm/llm_builder.py

 from typing import Union, Optional, List
 
 from langchain.callbacks.base import BaseCallbackHandler
-from langchain.llms.fake import FakeListLLM
 
 from core.constant import llm_constant
 from core.llm.error import ProviderTokenNotInitError
@@ -32,10 +31,7 @@ class LLMBuilder:
     """
 
     @classmethod
-    def to_llm(cls, tenant_id: str, model_name: str, **kwargs) -> Union[StreamableOpenAI, StreamableChatOpenAI, FakeListLLM]:
-        if model_name == 'fake':
-            return FakeListLLM(responses=[])
-
+    def to_llm(cls, tenant_id: str, model_name: str, **kwargs) -> Union[StreamableOpenAI, StreamableChatOpenAI]:
         provider = cls.get_default_provider(tenant_id)
 
         mode = cls.get_mode_by_model(model_name)

api/core/llm/streamable_azure_chat_open_ai.py

-from langchain.callbacks.manager import CallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun
+from langchain.callbacks.manager import CallbackManagerForLLMRun, AsyncCallbackManagerForLLMRun, Callbacks
 from langchain.schema import BaseMessage, ChatResult, LLMResult
 from langchain.chat_models import AzureChatOpenAI
 from typing import Optional, List, Dict, Any
@@ -69,68 +69,22 @@ class StreamableAzureChatOpenAI(AzureChatOpenAI):
 
         return message_tokens
 
-    def _generate(
-            self,
-            messages: List[BaseMessage],
-            stop: Optional[List[str]] = None,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-            **kwargs: Any,
-    ) -> ChatResult:
-        self.callback_manager.on_llm_start(
-            {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages],
-            verbose=self.verbose
-        )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-        self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
-    async def _agenerate(
+    @handle_llm_exceptions
+    def generate(
             self,
-            messages: List[BaseMessage],
+            messages: List[List[BaseMessage]],
             stop: Optional[List[str]] = None,
-            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+            callbacks: Callbacks = None,
             **kwargs: Any,
-    ) -> ChatResult:
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages],
-                verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages],
-                verbose=self.verbose
-            )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_end(result, verbose=self.verbose)
-        else:
-            self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
-    @handle_llm_exceptions
-    def generate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
     ) -> LLMResult:
-        return super().generate(messages, stop)
+        return super().generate(messages, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
+            self,
+            messages: List[List[BaseMessage]],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(messages, stop)
+        return await super().agenerate(messages, stop, callbacks, **kwargs)

api/core/llm/streamable_azure_open_ai.py

-import os
-
+from langchain.callbacks.manager import Callbacks
 from langchain.llms import AzureOpenAI
 from langchain.schema import LLMResult
 from typing import Optional, List, Dict, Mapping, Any
@@ -53,12 +52,20 @@ class StreamableAzureOpenAI(AzureOpenAI):
 
     @handle_llm_exceptions
     def generate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return super().generate(prompts, stop)
+        return super().generate(prompts, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(prompts, stop)
+        return await super().agenerate(prompts, stop, callbacks, **kwargs)

api/core/llm/streamable_chat_open_ai.py

 import os
 
-from langchain.callbacks.manager import CallbackManagerForLLMRun
-from langchain.schema import BaseMessage, ChatResult, LLMResult
+from langchain.callbacks.manager import Callbacks
+from langchain.schema import BaseMessage, LLMResult
 from langchain.chat_models import ChatOpenAI
 from typing import Optional, List, Dict, Any
 
@@ -71,65 +71,22 @@ class StreamableChatOpenAI(ChatOpenAI):
 
         return message_tokens
 
-    def _generate(
-            self,
-            messages: List[BaseMessage],
-            stop: Optional[List[str]] = None,
-            run_manager: Optional[CallbackManagerForLLMRun] = None,
-            **kwargs: Any,
-    ) -> ChatResult:
-        self.callback_manager.on_llm_start(
-            {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages], verbose=self.verbose
-        )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-        self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
-    async def _agenerate(
+    @handle_llm_exceptions
+    def generate(
             self,
-            messages: List[BaseMessage],
+            messages: List[List[BaseMessage]],
             stop: Optional[List[str]] = None,
-            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+            callbacks: Callbacks = None,
             **kwargs: Any,
-    ) -> ChatResult:
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages], verbose=self.verbose
-            )
-        else:
-            self.callback_manager.on_llm_start(
-                {"name": self.__class__.__name__}, [(message.type + ": " + message.content) for message in messages], verbose=self.verbose
-            )
-
-        chat_result = super()._generate(messages, stop)
-
-        result = LLMResult(
-            generations=[chat_result.generations],
-            llm_output=chat_result.llm_output
-        )
-
-        if self.callback_manager.is_async:
-            await self.callback_manager.on_llm_end(result, verbose=self.verbose)
-        else:
-            self.callback_manager.on_llm_end(result, verbose=self.verbose)
-
-        return chat_result
-
-    @handle_llm_exceptions
-    def generate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
     ) -> LLMResult:
-        return super().generate(messages, stop)
+        return super().generate(messages, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, messages: List[List[BaseMessage]], stop: Optional[List[str]] = None
+            self,
+            messages: List[List[BaseMessage]],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(messages, stop)
+        return await super().agenerate(messages, stop, callbacks, **kwargs)

api/core/llm/streamable_open_ai.py

 import os
 
+from langchain.callbacks.manager import Callbacks
 from langchain.schema import LLMResult
 from typing import Optional, List, Dict, Any, Mapping
 from langchain import OpenAI
@@ -48,15 +49,22 @@ class StreamableOpenAI(OpenAI):
             "organization": self.openai_organization if self.openai_organization else None,
         }}
 
-
     @handle_llm_exceptions
     def generate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return super().generate(prompts, stop)
+        return super().generate(prompts, stop, callbacks, **kwargs)
 
     @handle_llm_exceptions_async
     async def agenerate(
-            self, prompts: List[str], stop: Optional[List[str]] = None
+            self,
+            prompts: List[str],
+            stop: Optional[List[str]] = None,
+            callbacks: Callbacks = None,
+            **kwargs: Any,
     ) -> LLMResult:
-        return await super().agenerate(prompts, stop)
+        return await super().agenerate(prompts, stop, callbacks, **kwargs)

api/core/memory/read_only_conversation_token_db_string_buffer_shared_memory.py

 from typing import Any, List, Dict
 
 from langchain.memory.chat_memory import BaseChatMemory
-from langchain.schema import get_buffer_string, BaseMessage, BaseLanguageModel
+from langchain.schema import get_buffer_string
 
 from core.memory.read_only_conversation_token_db_buffer_shared_memory import \
     ReadOnlyConversationTokenDBBufferSharedMemory

api/extensions/ext_vector_store.py

-from core.vector_store.vector_store import VectorStore
-
-vector_store = VectorStore()
-
-
-def init_app(app):
-    vector_store.init_app(app)