Skip to content

Commit

Permalink
move NewTokenHandler callaback in itw own file
Browse files Browse the repository at this point in the history
  • Loading branch information
pieroit committed Oct 17, 2023
1 parent 5309f50 commit c724b99
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 11 deletions.
17 changes: 6 additions & 11 deletions core/cat/looking_glass/agent_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,24 +6,15 @@
from langchain.docstore.document import Document
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.callbacks.base import BaseCallbackHandler
from langchain.agents import AgentExecutor, LLMSingleActionAgent

from cat.looking_glass import prompts
from cat.looking_glass.callbacks import NewTokenHandler
from cat.looking_glass.output_parser import ToolOutputParser
from cat.utils import verbal_timedelta
from cat.log import log


class NewTokenHandler(BaseCallbackHandler):

def __init__(self, cat):
self.cat = cat
def on_llm_new_token(self, token: str, **kwargs) -> None:
self.cat.send_ws_message(token, "chat_token")



class AgentManager:
"""Manager of Langchain Agent.
Expand Down Expand Up @@ -54,7 +45,11 @@ def execute_tool_agent(self, agent_input, allowed_tools):
)

# main chain
agent_chain = LLMChain(prompt=prompt, llm=self.cat._llm, verbose=True)
agent_chain = LLMChain(
prompt=prompt,
llm=self.cat._llm,
verbose=True
)

# init agent
agent = LLMSingleActionAgent(
Expand Down
11 changes: 11 additions & 0 deletions core/cat/looking_glass/callbacks.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@

from langchain.callbacks.base import BaseCallbackHandler


class NewTokenHandler(BaseCallbackHandler):

def __init__(self, cat):
self.cat = cat

def on_llm_new_token(self, token: str, **kwargs) -> None:
self.cat.send_ws_message(token, "chat_token")

0 comments on commit c724b99

Please sign in to comment.