From 075de91675ebb015cc1dc2385d4d5ee88f8a23bd Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Sun, 19 Mar 2023 16:42:48 -0700
Subject: [PATCH] dbpedia

---
 langchain/chains/dbpedia/__init__.py |  0
 langchain/chains/dbpedia/base.py     | 61 ++++++++++++++++++++++++++++
 langchain/chains/dbpedia/prompt.py   | 52 ++++++++++++++++++++++++
 3 files changed, 113 insertions(+)
 create mode 100644 langchain/chains/dbpedia/__init__.py
 create mode 100644 langchain/chains/dbpedia/base.py
 create mode 100644 langchain/chains/dbpedia/prompt.py

diff --git a/langchain/chains/dbpedia/__init__.py b/langchain/chains/dbpedia/__init__.py
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/langchain/chains/dbpedia/base.py b/langchain/chains/dbpedia/base.py
new file mode 100644
index 00000000000..67431c9dbe2
--- /dev/null
+++ b/langchain/chains/dbpedia/base.py
@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from langchain.chains.base import Chain
+from langchain.chains.dbpedia.prompt import ANSWER_PROMPT_SELECTOR, PROMPT_SELECTOR
+from langchain.chains.llm import LLMChain
+from langchain.prompts.base import BasePromptTemplate
+from langchain.schema import BaseLanguageModel
+
+
+class DBPediaChain(Chain):
+    query_chain: LLMChain
+    answer_chain: LLMChain
+    input_key: str = "question"
+    output_key: str = "answer"
+
+    @classmethod
+    def from_llm(
+        cls,
+        llm: BaseLanguageModel,
+        query_prompt: Optional[BasePromptTemplate] = None,
+        answer_prompt: Optional[BasePromptTemplate] = None,
+        **kwargs: Any,
+    ) -> DBPediaChain:
+        query_prompt = query_prompt or PROMPT_SELECTOR.get_prompt(llm)
+        query_chain = LLMChain(llm=llm, prompt=query_prompt)
+        answer_prompt = answer_prompt or ANSWER_PROMPT_SELECTOR.get_prompt(llm)
+        answer_chain = LLMChain(llm=llm, prompt=answer_prompt)
+        return cls(query_chain=query_chain, answer_chain=answer_chain, **kwargs)
+
+    @property
+    def input_keys(self) -> List[str]:
+        return [self.input_key]
+
+    @property
+    def output_keys(self) -> List[str]:
+        return [self.output_key]
+
+    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
+        from SPARQLWrapper import JSON, SPARQLWrapper
+
+        sparql = SPARQLWrapper("http://dbpedia.org/sparql")
+        sparql.setReturnFormat(JSON)
+        query = self.query_chain.run(inputs[self.input_key])
+        self.callback_manager.on_text("Query written:", end="\n", verbose=self.verbose)
+        self.callback_manager.on_text(
+            query, color="green", end="\n", verbose=self.verbose
+        )
+        sparql.setQuery(query)
+        result = sparql.query().convert()
+        self.callback_manager.on_text(
+            "Response gotten:", end="\n", verbose=self.verbose
+        )
+        self.callback_manager.on_text(
+            result, color="green", end="\n", verbose=self.verbose
+        )
+        answer = self.answer_chain.run(
+            question=inputs[self.input_key], query=query, response=result
+        )
+        return {self.output_key: answer}
diff --git a/langchain/chains/dbpedia/prompt.py b/langchain/chains/dbpedia/prompt.py
new file mode 100644
index 00000000000..a8056d3a083
--- /dev/null
+++ b/langchain/chains/dbpedia/prompt.py
@@ -0,0 +1,52 @@
+from langchain.chains.prompt_selector import ConditionalPromptSelector, is_chat_model
+from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
+from langchain.prompts.prompt import PromptTemplate
+from langchain.schema import HumanMessage
+
+TEMPLATE = """Write a sparkql query to execute against DBPedia to answer the following question
+
+Question: {question}
+SPARQL Query:"""
+PROMPT = PromptTemplate.from_template(TEMPLATE)
+
+INSTRUCTIONS_TEMPLATE = """Write a sparkql query to execute against DBPedia to answer the following question.
+Your answer should be a valid SPARKQL query and NOTHING else.
+Always return just a SPARKQL query."""
+INSTRUCTIONS = HumanMessage(content=INSTRUCTIONS_TEMPLATE)
+CHAT_PROMPT = ChatPromptTemplate.from_messages(
+    [INSTRUCTIONS, HumanMessagePromptTemplate.from_template("{question}")]
+)
+
+PROMPT_SELECTOR = ConditionalPromptSelector(
+    default_prompt=PROMPT, conditionals=[(is_chat_model, CHAT_PROMPT)]
+)
+
+ANSWER_TEMPLATE = """Write a sparkql query to execute against DBPedia to answer the following question
+
+Question: {question}
+SPARKQL Query: {query}
+SPARKQL Response: {response}
+Final Answer (in plain English):"""
+ANSWER_PROMPT = PromptTemplate.from_template(ANSWER_TEMPLATE)
+
+ANSWER_INSTRUCTIONS_TEMPLATE = """I wrote this SPARKQL query:
+----------
+{query}
+----------
+
+I got this response:
+----------
+{response}
+----------
+
+Now, use the above information to answer my next question."""
+ANSWER_INSTRUCTIONS = HumanMessagePromptTemplate.from_template(
+    ANSWER_INSTRUCTIONS_TEMPLATE
+)
+ANSWER_CHAT_PROMPT = ChatPromptTemplate.from_messages(
+    [ANSWER_INSTRUCTIONS, HumanMessagePromptTemplate.from_template("{question}")]
+)
+
+ANSWER_PROMPT_SELECTOR = ConditionalPromptSelector(
+    default_prompt=ANSWER_PROMPT, conditionals=[(is_chat_model, ANSWER_CHAT_PROMPT)]
+)