Add input types to cypher templates (langchain-ai#12800)

ElveStudio · Nov 2, 2023 · 2a9f40e · 2a9f40e
1 parent c4fdf78
commit 2a9f40e
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 5 deletions.
diff --git a/templates/neo4j-cypher-ft/neo4j_cypher_ft/chain.py b/templates/neo4j-cypher-ft/neo4j_cypher_ft/chain.py
@@ -5,14 +5,10 @@
 from langchain.chat_models import ChatOpenAI
 from langchain.graphs import Neo4jGraph
 from langchain.prompts import ChatPromptTemplate
+from langchain.pydantic_v1 import BaseModel, Field
 from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import RunnablePassthrough
 
-try:
-    from pydantic.v1.main import BaseModel, Field
-except ImportError:
-    from pydantic.main import BaseModel, Field
-
 # Connection to Neo4j
 graph = Neo4jGraph()
 
@@ -127,3 +123,12 @@ def map_to_database(entities: Entities) -> Optional[str]:
     | qa_llm
     | StrOutputParser()
 )
+
+# Add typing for input
+
+
+class Question(BaseModel):
+    question: str
+
+
+chain = chain.with_types(input_type=Question)
diff --git a/templates/neo4j-cypher/neo4j_cypher/chain.py b/templates/neo4j-cypher/neo4j_cypher/chain.py
@@ -2,6 +2,7 @@
 from langchain.chat_models import ChatOpenAI
 from langchain.graphs import Neo4jGraph
 from langchain.prompts import ChatPromptTemplate
+from langchain.pydantic_v1 import BaseModel
 from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import RunnablePassthrough
 
@@ -71,3 +72,12 @@
     | qa_llm
     | StrOutputParser()
 )
+
+# Add typing for input
+
+
+class Question(BaseModel):
+    question: str
+
+
+chain = chain.with_types(input_type=Question)