cli[minor]: Improve partner migrations (#20938)

This auto generates partner migrations.

At the moment the migration is from community -> partner.

So one would need to run the migration script twice to go from langchain to partner.
This commit is contained in:
Eugene Yurtsev
2024-04-26 12:30:15 -04:00
committed by GitHub
parent 5653f36adc
commit 12c906f6ce
17 changed files with 265 additions and 50 deletions

View File

@@ -1,7 +1,7 @@
from tests.unit_tests.migrate.integration.case import Case
from tests.unit_tests.migrate.integration.cases import imports
from tests.unit_tests.migrate.integration.file import File
from tests.unit_tests.migrate.integration.folder import Folder
from tests.unit_tests.migrate.cli_runner.case import Case
from tests.unit_tests.migrate.cli_runner.cases import imports
from tests.unit_tests.migrate.cli_runner.file import File
from tests.unit_tests.migrate.cli_runner.folder import Folder
cases = [
Case(

View File

@@ -1,5 +1,5 @@
from tests.unit_tests.migrate.integration.case import Case
from tests.unit_tests.migrate.integration.file import File
from tests.unit_tests.migrate.cli_runner.case import Case
from tests.unit_tests.migrate.cli_runner.file import File
cases = [
Case(
@@ -7,7 +7,7 @@ cases = [
source=File(
"app.py",
content=[
"from langchain.chat_models import ChatOpenAI",
"from langchain_community.chat_models import ChatOpenAI",
"",
"",
"class foo:",

View File

@@ -28,4 +28,19 @@ def test_generate_migrations() -> None:
"langchain_community.chat_models.azure_openai.AzureChatOpenAI",
"langchain_openai.AzureChatOpenAI",
),
("langchain_community.llms.AzureOpenAI", "langchain_openai.AzureOpenAI"),
("langchain_community.llms.OpenAI", "langchain_openai.OpenAI"),
(
"langchain_community.embeddings.AzureOpenAIEmbeddings",
"langchain_openai.AzureOpenAIEmbeddings",
),
(
"langchain_community.embeddings.OpenAIEmbeddings",
"langchain_openai.OpenAIEmbeddings",
),
(
"langchain_community.chat_models.AzureChatOpenAI",
"langchain_openai.AzureChatOpenAI",
),
("langchain_community.chat_models.ChatOpenAI", "langchain_openai.ChatOpenAI"),
]

View File

@@ -19,22 +19,32 @@ class TestReplaceImportsCommand(CodemodTest):
from langchain.chat_models import ChatOpenAI
"""
after = """
from langchain_community.chat_models import ChatOpenAI
"""
self.assertCodemod(before, after)
def test_from_community_to_partner(self) -> None:
"""Test that we can replace imports from community to partner."""
before = """
from langchain_community.chat_models import ChatOpenAI
"""
after = """
from langchain_openai import ChatOpenAI
"""
self.assertCodemod(before, after)
def test_noop_import(self) -> None:
code = """
from foo import ChatOpenAI
from foo import ChatOpenAI
"""
self.assertCodemod(code, code)
def test_mixed_imports(self) -> None:
before = """
from langchain.chat_models import ChatOpenAI, ChatAnthropic, foo
from langchain_community.chat_models import ChatOpenAI, ChatAnthropic, foo
"""
after = """
from langchain.chat_models import foo
from langchain_community.chat_models import foo
from langchain_anthropic import ChatAnthropic
from langchain_openai import ChatOpenAI
"""