astradb

app.py

@@ -8,8 +8,12 @@ def run(text):
     now = datetime.datetime.now(pytz.timezone('Asia/Tokyo'))
     print('\ntext:',text,now)
     result = pipe.run({
-        "prompt_builder":{"text":text}
-    })
+        'retrieve_prompt_builder':{'text':text},
+        'prompt_builder':{'text':text}
+    },include_outputs_from=['embedding_retriever','parser'])
+    for document in result['embedding_retriever']['documents']:
+        print(document.score,document.id)
+        print(document.content)
     parsed = result['parser']['parsed']
     print(parsed)
     return json.dumps(parsed,indent=2,ensure_ascii=False)

pipe.py

@@ -1,18 +1,31 @@
-import json, json_repair
-from typing import List, Dict
-from haystack import Pipeline, component
+from haystack import Document, Pipeline,component
+from haystack.components.embedders import SentenceTransformersDocumentEmbedder,SentenceTransformersTextEmbedder
+from haystack_integrations.document_stores.astra import AstraDocumentStore
+from haystack.components.writers import DocumentWriter
+from haystack.document_stores.types import DuplicatePolicy
 from haystack.components.builders import PromptBuilder
 from haystack_integrations.components.generators.google_ai import GoogleAIGeminiGenerator
+from haystack_integrations.components.retrievers.astra import AstraEmbeddingRetriever
+import json, json_repair
+from typing import List, Dict
+
+document_store = AstraDocumentStore(
+    collection_name='company_contents',
+    duplicates_policy=DuplicatePolicy.OVERWRITE
+)
+
+retrieve_template = '''
+    次の日本語の文章の明確さを判定したい
+    テキスト: {{text}}
+'''
 
 template = """
-    次のテキストに基づいて、JSON 形式で 回答を生成して下さい。
-    テキストの「明確さ」を様々な点から評価し、100点満点で点数を採点してください。
+    どのように日本語の文章を記載するとよいかまとめた自社の資料があります。
+    この資料を参考にして、以下のテキストの「明確さ」を様々な点から評価し、100点満点で点数を採点してください。
     次にその点数とした理由の解説と、どの辺が曖昧だったかを説明してください。
     テキストをより明確にする上で修正するべき箇所をピックアップし、その箇所が何行目の何文字目であるかと、修正前のテキストと修正後のテキストをまとめて提案して下さい。
-
     以上のことをマークダウンや説明なしで、JSON のみで回答して下さい。
     次のようなJSONフォーマット形式に必ず従ってください。
-
     {
         "score" : 採点した点数(数値) ,
         "description": "その点数になった理由",
@@ -24,7 +37,10 @@ template = """
             }, ...
         ]
     }
-
+    資料：
+      {% for document in documents %}
+        {{ document.content }}
+      {% endfor %}
     テキスト: {{text}}
 """
 
@@ -45,13 +61,23 @@ class ResultParser:
             parsed = parsed[0]
         return {"parsed": parsed}
 
+retrieve_prompt_builder = PromptBuilder(template=retrieve_template)
+text_embedder =  SentenceTransformersTextEmbedder(model=os.environ['embedding_model'])
+embedding_retriever = AstraEmbeddingRetriever(document_store=document_store)
 prompt_builder =  PromptBuilder(template=template)
-llm = GoogleAIGeminiGenerator(model="models/gemini-1.5-flash")
+llm = GoogleAIGeminiGenerator(model=os.environ["gemini_model"])
 parser = ResultParser()
+text_embedder.warm_up()
 
 pipe = Pipeline()
+pipe.add_component("retrieve_prompt_builder",retrieve_prompt_builder)
+pipe.add_component("text_embedder",text_embedder)
+pipe.add_component("embedding_retriever",embedding_retriever)
 pipe.add_component("prompt_builder",prompt_builder)
 pipe.add_component("llm",llm)
 pipe.add_component("parser",parser)
+pipe.connect("retrieve_prompt_builder.prompt", "text_embedder.text")
+pipe.connect("text_embedder", "embedding_retriever")
+pipe.connect("embedding_retriever.documents", "prompt_builder.documents")
 pipe.connect("prompt_builder", "llm")
 pipe.connect("llm.replies", "parser.replies")

requirements.txt

@@ -5,4 +5,5 @@ haystack-ai
 google-ai-haystack
 accelerate
 pytz
-json_repair
+json_repair
+astra-haystack