forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.3 KiB
Python
41 lines
1.3 KiB
Python
from abc import abstractmethod
|
|
from typing import Any, List, Optional
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from langchain.agents.plan_and_execute.schema import Plan, PlanOutputParser
|
|
from langchain.callbacks.manager import Callbacks
|
|
from langchain.chains.llm import LLMChain
|
|
|
|
|
|
class BasePlanner(BaseModel):
|
|
@abstractmethod
|
|
def plan(self, inputs: dict, callbacks: Callbacks = None, **kwargs: Any) -> Plan:
|
|
"""Given input, decided what to do."""
|
|
|
|
@abstractmethod
|
|
async def aplan(
|
|
self, inputs: dict, callbacks: Callbacks = None, **kwargs: Any
|
|
) -> Plan:
|
|
"""Given input, decided what to do."""
|
|
|
|
|
|
class LLMPlanner(BasePlanner):
|
|
llm_chain: LLMChain
|
|
output_parser: PlanOutputParser
|
|
stop: Optional[List] = None
|
|
|
|
def plan(self, inputs: dict, callbacks: Callbacks = None, **kwargs: Any) -> Plan:
|
|
"""Given input, decided what to do."""
|
|
llm_response = self.llm_chain.run(**inputs, stop=self.stop, callbacks=callbacks)
|
|
return self.output_parser.parse(llm_response)
|
|
|
|
async def aplan(
|
|
self, inputs: dict, callbacks: Callbacks = None, **kwargs: Any
|
|
) -> Plan:
|
|
"""Given input, decided what to do."""
|
|
llm_response = await self.llm_chain.arun(
|
|
**inputs, stop=self.stop, callbacks=callbacks
|
|
)
|
|
return self.output_parser.parse(llm_response)
|