jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

🔨 refactor(conftest.py): reformat code for better readability and maintainability

Browse source 
✨ feat(conftest.py): add MyCustomChain class as an example of a custom chain
✨ feat(conftest.py): add CustomChain class as a custom component for building a document
✨ feat(conftest.py): add CSVLoaderComponent class as a custom component for loading CSV files and converting rows to documents
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-07-14 14:05:57 -03:00
commit 99ef7c728d
1 changed files with 120 additions and 90 deletions
Download patch file Download diff file Expand all files Collapse all files

210
tests/conftest.py
View file

@ -122,116 +122,146 @@ def client_fixture(session: Session): #
def custom_chain(): def custom_chain():
return '''from __future__ import annotations return '''from __future__ import annotations
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from pydantic import Extra from pydantic import Extra
from langchain.schema import BaseLanguageModel, Document from langchain.schema import BaseLanguageModel, Document
from langchain.callbacks.manager import ( from langchain.callbacks.manager import (
AsyncCallbackManagerForChainRun, AsyncCallbackManagerForChainRun,
CallbackManagerForChainRun, CallbackManagerForChainRun,
) )
from langchain.chains.base import Chain from langchain.chains.base import Chain
from langchain.prompts import StringPromptTemplate from langchain.prompts import StringPromptTemplate
from langflow.interface.custom.base import CustomComponent from langflow.interface.custom.base import CustomComponent
class MyCustomChain(Chain): class MyCustomChain(Chain):
"""
An example of a custom chain.
"""
prompt: StringPromptTemplate
"""Prompt object to use."""
llm: BaseLanguageModel
output_key: str = "text" #: :meta private:
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
arbitrary_types_allowed = True
@property
def input_keys(self) -> List[str]:
"""Will be whatever keys the prompt expects.
:meta private:
""" """
An example of a custom chain. return self.prompt.input_variables
@property
def output_keys(self) -> List[str]:
"""Will always return text key.
:meta private:
""" """
return [self.output_key]
prompt: StringPromptTemplate def _call(
"""Prompt object to use.""" self,
llm: BaseLanguageModel inputs: Dict[str, Any],
output_key: str = "text" #: :meta private: run_manager: Optional[CallbackManagerForChainRun] = None,
) -> Dict[str, str]:
# Your custom chain logic goes here
# This is just an example that mimics LLMChain
prompt_value = self.prompt.format_prompt(**inputs)
class Config: # Whenever you call a language model, or another chain, you should pass
"""Configuration for this pydantic object.""" # a callback manager to it. This allows the inner run to be tracked by
# any callbacks that are registered on the outer run.
# You can always obtain a callback manager for this by calling
# `run_manager.get_child()` as shown below.
response = self.llm.generate_prompt(
[prompt_value],
callbacks=run_manager.get_child() if run_manager else None,
)
extra = Extra.forbid # If you want to log something about this run, you can do so by calling
arbitrary_types_allowed = True # methods on the `run_manager`, as shown below. This will trigger any
# callbacks that are registered for that event.
if run_manager:
run_manager.on_text("Log something about this run")
@property return {self.output_key: response.generations[0][0].text}
def input_keys(self) -> List[str]:
"""Will be whatever keys the prompt expects.
:meta private: async def _acall(
""" self,
return self.prompt.input_variables inputs: Dict[str, Any],
run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
) -> Dict[str, str]:
# Your custom chain logic goes here
# This is just an example that mimics LLMChain
prompt_value = self.prompt.format_prompt(**inputs)
@property # Whenever you call a language model, or another chain, you should pass
def output_keys(self) -> List[str]: # a callback manager to it. This allows the inner run to be tracked by
"""Will always return text key. # any callbacks that are registered on the outer run.
# You can always obtain a callback manager for this by calling
# `run_manager.get_child()` as shown below.
response = await self.llm.agenerate_prompt(
[prompt_value],
callbacks=run_manager.get_child() if run_manager else None,
)
:meta private: # If you want to log something about this run, you can do so by calling
""" # methods on the `run_manager`, as shown below. This will trigger any
return [self.output_key] # callbacks that are registered for that event.
if run_manager:
await run_manager.on_text("Log something about this run")
def _call( return {self.output_key: response.generations[0][0].text}
self,
inputs: Dict[str, Any],
run_manager: Optional[CallbackManagerForChainRun] = None,
) -> Dict[str, str]:
# Your custom chain logic goes here
# This is just an example that mimics LLMChain
prompt_value = self.prompt.format_prompt(**inputs)
# Whenever you call a language model, or another chain, you should pass @property
# a callback manager to it. This allows the inner run to be tracked by def _chain_type(self) -> str:
# any callbacks that are registered on the outer run. return "my_custom_chain"
# You can always obtain a callback manager for this by calling
# `run_manager.get_child()` as shown below.
response = self.llm.generate_prompt(
[prompt_value],
callbacks=run_manager.get_child() if run_manager else None,
)
# If you want to log something about this run, you can do so by calling class CustomChain(CustomComponent):
# methods on the `run_manager`, as shown below. This will trigger any display_name: str = "Custom Chain"
# callbacks that are registered for that event. field_config = {
if run_manager: "prompt": {"field_type": "prompt"},
run_manager.on_text("Log something about this run") "llm": {"field_type": "BaseLanguageModel"},
}
return {self.output_key: response.generations[0][0].text} def build(self, prompt, llm, input: str) -> Document:
chain = MyCustomChain(prompt=prompt, llm=llm)
return chain(input)'''
async def _acall(
self,
inputs: Dict[str, Any],
run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
) -> Dict[str, str]:
# Your custom chain logic goes here
# This is just an example that mimics LLMChain
prompt_value = self.prompt.format_prompt(**inputs)
# Whenever you call a language model, or another chain, you should pass @pytest.fixture
# a callback manager to it. This allows the inner run to be tracked by def data_processing():
# any callbacks that are registered on the outer run. return """import pandas as pd
# You can always obtain a callback manager for this by calling from langchain.schema import Document
# `run_manager.get_child()` as shown below. from langflow.interface.custom.base import CustomComponent
response = await self.llm.agenerate_prompt(
[prompt_value],
callbacks=run_manager.get_child() if run_manager else None,
)
# If you want to log something about this run, you can do so by calling class CSVLoaderComponent(CustomComponent):
# methods on the `run_manager`, as shown below. This will trigger any display_name: str = "CSV Loader"
# callbacks that are registered for that event. field_config = {
if run_manager: "filename": {"field_type": "str", "required": True},
await run_manager.on_text("Log something about this run") "column_name": {"field_type": "str", "required": True},
}
return {self.output_key: response.generations[0][0].text} def build(self, filename: str, column_name: str) -> List[Document]:
# Load the CSV file
df = pd.read_csv(filename)
@property # Verify the column exists
def _chain_type(self) -> str: if column_name not in df.columns:
return "my_custom_chain" raise ValueError(f"Column '{column_name}' not found in the CSV file")
class CustomChain(CustomComponent): # Convert each row of the specified column to a document object
display_name: str = "Custom Chain" documents = []
field_config = { for content in df[column_name]:
"prompt": {"field_type": "prompt"}, metadata = {"filename": filename}
"llm": {"field_type": "BaseLanguageModel"}, documents.append(Document(page_content=str(content), metadata=metadata))
}
def build(self, prompt, llm, input: str) -> Document: return documents"""
chain = MyCustomChain(prompt=prompt, llm=llm)
return chain(input)'''