use StrOutputParser in /fix

jupyterlab · Jan 16, 2025 · c6cfe8f · c6cfe8f
1 parent b69e2f7
commit c6cfe8f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py
@@ -3,6 +3,7 @@
 from jupyter_ai.models import CellWithErrorSelection, HumanChatMessage
 from jupyter_ai_magics.providers import BaseProvider
 from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
 
 from .base import BaseChatHandler, SlashCommandRoutingType
 
@@ -76,7 +77,7 @@ def create_llm_chain(
         self.llm = llm
         prompt_template = FIX_PROMPT_TEMPLATE
 
-        runnable = prompt_template | llm  # type:ignore
+        runnable = prompt_template | llm | StrOutputParser()  # type:ignore
         self.llm_chain = runnable
 
     async def process_message(self, message: HumanChatMessage):