Update ZeroShotAgent and OpenAIComponent

This commit is contained in:
anovazzi1 2024-01-15 20:25:43 -03:00
commit fbd64033dc
2 changed files with 6 additions and 6 deletions

View file

@ -3,9 +3,7 @@ from langflow import CustomComponent
from langchain.agents import ZeroShotAgent from langchain.agents import ZeroShotAgent
from langchain_core.tools import BaseTool from langchain_core.tools import BaseTool
from typing import List, Optional from typing import List, Optional
from langflow.field_typing import ( from langflow.components.chains.LLMChain import LLMChain
BaseLanguageModel,
)
class ZeroShotAgentComponent(CustomComponent): class ZeroShotAgentComponent(CustomComponent):
display_name = "ZeroShotAgent" display_name = "ZeroShotAgent"
@ -21,9 +19,9 @@ class ZeroShotAgentComponent(CustomComponent):
def build( def build(
self, self,
llm: BaseLanguageModel, llm: LLMChain,
tools: List[BaseTool], tools: List[BaseTool],
prefix: Optional[str] = "Answer the following questions as best you can. You have access to the following tools:", prefix: Optional[str] = "Answer the following questions as best you can. You have access to the following tools:",
suffix: Optional[str] = "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}", suffix: Optional[str] = "Begin!\n\nQuestion: {input}\nThought:{agent_scratchpad}",
) -> ZeroShotAgent: ) -> ZeroShotAgent:
return ZeroShotAgent(llm=llm, tools=tools, prefix=prefix, suffix=suffix) return ZeroShotAgent(llm_chain=llm, tools=tools, prefix=prefix, suffix=suffix)

View file

@ -42,10 +42,12 @@ class OpenAIComponent(CustomComponent):
max_tokens: Optional[int] = 256, max_tokens: Optional[int] = 256,
model_kwargs: Optional[Dict] = None, model_kwargs: Optional[Dict] = None,
model_name: Optional[str] = "text-davinci-003", model_name: Optional[str] = "text-davinci-003",
openai_api_base: Optional[str] = "https://api.openai.com/v1", openai_api_base: Optional[str] = "",
openai_api_key: str = "", openai_api_key: str = "",
temperature: Optional[float] = 0.7, temperature: Optional[float] = 0.7,
) -> OpenAI: ) -> OpenAI:
if(not openai_api_base):
openai_api_base = "https://api.openai.com/v1"
return OpenAI( return OpenAI(
max_tokens=max_tokens, max_tokens=max_tokens,
model_kwargs=model_kwargs or {}, model_kwargs=model_kwargs or {},