| 1 | from typing import Callable, TypedDict |
| 2 | from langchain.prompts import ( |
| 3 | ChatPromptTemplate, |
| 4 | FewShotChatMessagePromptTemplate, |
| 5 | ) |
| 6 | |
| 7 | from langchain.schema import AIMessage |
| 8 | from langchain_core.messages import HumanMessage, SystemMessage |
| 9 | |
| 10 | from langchain_core.language_models.chat_models import BaseChatModel |
| 11 | from langchain_core.language_models.llms import BaseLLM |
| 12 | |
| 13 | |
| 14 | class Example(TypedDict): |
| 15 | input: str |
| 16 | output: str |
| 17 | |
| 18 | async def call_llm( |
| 19 | system: str, |
| 20 | model: BaseChatModel | BaseLLM, |
| 21 | message: str, |
| 22 | examples: list[Example] = [], |
| 23 | callback: Callable[[str], None] | None = None |
| 24 | ): |
| 25 | |
| 26 | example_prompt = ChatPromptTemplate.from_messages( |
| 27 | [ |
| 28 | HumanMessage(content="{input}"), |
| 29 | AIMessage(content="{output}"), |
| 30 | ] |
| 31 | ) |
| 32 | |
| 33 | few_shot_prompt = FewShotChatMessagePromptTemplate( |
| 34 | example_prompt=example_prompt, |
| 35 | examples=examples, # type: ignore |
| 36 | input_variables=[], |
| 37 | ) |
| 38 | |
| 39 | few_shot_prompt.format() |
| 40 | |
| 41 | |
| 42 | final_prompt = ChatPromptTemplate.from_messages( |
| 43 | [ |
| 44 | SystemMessage(content=system), |
| 45 | few_shot_prompt, |
| 46 | HumanMessage(content=message), |
| 47 | ] |
| 48 | ) |
| 49 | |
| 50 | chain = final_prompt | model |
| 51 | |
| 52 | response = "" |
| 53 | async for chunk in chain.astream({}): |
| 54 | # await self.handle_intervention() # wait for intervention and handle it, if paused |
| 55 | |
| 56 | if isinstance(chunk, str): |
| 57 | content = chunk |
| 58 | elif hasattr(chunk, "content"): |
| 59 | content = str(chunk.content) |
| 60 | else: |
| 61 | content = str(chunk) |
| 62 | |
| 63 | if callback: |
| 64 | callback(content) |
| 65 | |
| 66 | response += content |
| 67 | |
| 68 | return response |
| 69 |