Coverage for app/endpoint/knowledge/create_knowledge.py: 40%

1from app.endpoint.knowledge.router import router 

2from app.infrastructure.database.db import get_session 

3 

4from fastapi import Depends 

5from sqlmodel import Session 

6 

7# ========== 埋め込み関数 ========== 

8from ollama import Client 

9OLLAMA_URL = "http://host.docker.internal:11434" 

10OLLAMA_EMMBEDDING_MODEL = "zylonai/multilingual-e5-large:latest" 

11OLLAMA_LLM_MODEL = "llama3:latest" 

12def get_embedding(text: str): 

13 # ollamaクライアント 

14 client = Client(host=OLLAMA_URL) 

15 # emb実行 

16 result = client.embed(model=OLLAMA_EMMBEDDING_MODEL, input=text) 

17 return result 

18 

19# ========== 簡易グラフ変換関数 ========== 

20import json 

21def get_graph_by_manual(text: str): 

22 """ 

23 非LLMなルールベース処理。 

24 例文：「田中さんはABC株式会社でエンジニアをしています。」 

25 """ 

26 if "は" in text and "で" in text: 

27 person = text.split("は")[0].strip() 

28 company = text.split("は")[-1].split("で")[0].strip() 

29 role = text.split("で")[-1].replace("をしています", "").replace("として働いています", "").strip("。") 

30 return {"person": person, "company": company, "role": role} 

31 return None 

32 

33def get_graph_by_ollama(text: str): 

34 prompt = f""" 

35あなたは自然文を構造化するAIアシスタントです。 

36以下の日本語文を、グラフ構造としてノードとリレーションに分解してください。 

37 

38出力形式は JSON としてください。形式は以下の通りです。 

39**必ずこの形式で返してください。余計な文章や説明は一切入れないでください。** 

40 

41{{ 

42 "nodes": [{{"id": "名前", "label": "ラベル"}}], 

43 "relationships": [{{"from": "ノードID", "to": "ノードID", "type": "リレーションタイプ", "properties": {{}} }}] 

44}} 

45 

46入力文： 

47{text} 

48""" 

49 # ollamaクライアント 

50 client = Client(host=OLLAMA_URL) 

51 response = client.generate( 

52 model=OLLAMA_LLM_MODEL, 

53 prompt=prompt 

54 ) 

55 

56 print(response['response']) 

57 # 応答を JSON にパース 

58 try: 

59 result = json.loads(response['response']) 

60 

61 except json.JSONDecodeError: 

62 raise ValueError("Ollamaからの応答をJSONとして解析できませんでした:\n" + response['response']) 

63 

64 return result 

65 

66# neo4j登録 

67from neo4j import GraphDatabase 

68NEO4J_URI = "bolt://search:7687" 

69NEO4J_USER = "neo4j" 

70NEO4J_PASS = "password" 

71 

72def register_graph(graph_data, text, embedding): 

73 driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASS)) 

74 with driver.session() as session: 

75 for node in graph_data.get("nodes", []): 

76 props = { 

77 "id": node["id"], 

78 "text": text, 

79 "embedding": embedding, 

80 } 

81 label = node.get("label", "Entity") 

82 session.run( 

83 f""" 

84 MERGE (n:{label} {{id: $id}}) 

85 SET n.text = $text, 

86 n.embedding = $embedding 

87 """, 

88 **props 

89 ) 

90 

91 for rel in graph_data.get("relationships", []): 

92 session.run( 

93 f""" 

94 MATCH (a {{id: $_from}}) 

95 MATCH (b {{id: $to}}) 

96 MERGE (a)-[r:{rel["type"]}]->(b) 

97 SET r += $props 

98 """, 

99 _from=rel["from"], 

100 to=rel["to"], 

101 props=rel.get("properties", {}) 

102 ) 

103 

104from sqlmodel import SQLModel 

105class CreateKnowledgeRequest(SQLModel): 

106 text: str 

107 

108@router.post( 

109 "/knowledges" 

110) 

111async def create_knowledge( 

112 request: CreateKnowledgeRequest, 

113 session: Session = Depends(get_session), 

114): 

115 text = request.text 

116 print(text) 

117 embedding = get_embedding(text)['embeddings'][0] 

118 graph = get_graph_by_ollama(text) 

119 # person, company, role = graph["person"], graph["company"], graph["role"] 

120 

121 # 3. Neo4j にノード＋リレーション登録 

122 register_graph(graph, text, embedding) 

123 

124 

125 return { 

126 "graph": graph, 

127 "text": text, 

128 # "neo4j": { 

129 # "person": person, 

130 # "company": company, 

131 # "role": role 

132 # }, 

133 "embedding": embedding, 

134 }