10 分钟写一个 tidb-ai 机器人帮你解答 tidb 问题

$ mysql \      --comments -u '3grcM9DRGNroFfR.JW2bS3MJ' -p'!C06Ao41fcuqiHnLr4lycGU6FmKAIwFFC0QO' \      -h gateway01.us-west-2.prod.aws.tidbcloud.com -P 4000 \      --ssl-mode=VERIFY_IDENTITY --ssl-ca=/etc/pki/tls/certs/ca-bundle.crt

DROP DATABASE IF EXISTS chatdb;   CREATE DATABASE chatdb;   EXIT

$ pip install \      click==8.1.7 \      PyMySQL==1.1.0 \      SQLAlchemy==2.0.30 \      tidb-vector==0.0.9 \      pydantic==2.7.1 pydantic_core==2.18.2 \      dspy-ai==2.4.12 \      langchain-community==0.2.0 \      wikipedia==1.4.0 \      pyvis==0.3.1 \      openai==1.27.0 \      zhipuai==2.1.3

$ export ZHIPUAI_API_KEY=${ZHIPUAI_API_KEY} #把智普的apikey填这里

$ cat > build-graph.py <<'EOF'   import os   import pymysql   import dspy   import enum   import openai      from zhipuai import ZhipuAI   from pymysql import Connection   from pymysql.cursors import DictCursor   from dspy.functional import TypedPredictor   from pydantic import BaseModel, Field   from typing import Mapping, Any, Optional, List   from langchain_community.document_loaders import WikipediaLoader   from pyvis.network import Network   from IPython.display import HTML   from sqlalchemy import (       Column,       Integer,       String,       Text,       JSON,       ForeignKey,       BLOB,       Enum as SQLEnum,       DateTime,       URL,       create_engine,       or_,   )   from sqlalchemy.orm import relationship, Session, declarative_base, joinedload   from tidb_vector.sqlalchemy import VectorType   class Entity(BaseModel):       """List of entities extracted from the text to form the knowledge graph"""       name: str = Field(           description="Name of the entity, it should be a clear and concise term"       )       description: str = Field(           description=(               "Description of the entity, it should be a complete and comprehensive sentence, not few words. "               "Sample description of entity 'TiDB in-place upgrade': "               "'Upgrade TiDB component binary files to achieve upgrade, generally use rolling upgrade method'"           )       )   class Relationship(BaseModel):       """List of relationships extracted from the text to form the knowledge graph"""       source_entity: str = Field(           description="Source entity name of the relationship, it should an existing entity in the Entity list"       )       target_entity: str = Field(           description="Target entity name of the relationship, it should an existing entity in the Entity list"       )       relationship_desc: str = Field(           description=(               "Description of the relationship, it should be a complete and comprehensive sentence, not few words. "               "Sample relationship description: 'TiDB will release a new LTS version every 6 months.'"           )       )   class KnowledgeGraph(BaseModel):       """Graph representation of the knowledge for text."""       entities: List[Entity] = Field(           description="List of entities in the knowledge graph"       )       relationships: List[Relationship] = Field(           description="List of relationships in the knowledge graph"       )   class ExtractGraphTriplet(dspy.Signature):       text = dspy.InputField(           desc="a paragraph of text to extract entities and relationships to form a knowledge graph"       )       knowledge: KnowledgeGraph = dspy.OutputField(           desc="Graph representation of the knowledge extracted from the text."       )   class Extractor(dspy.Module):       def __init__(self):           super().__init__()           self.prog_graph = TypedPredictor(ExtractGraphTriplet)       def forward(self, text):           return self.prog_graph(               text=text,               config={                   "response_format": {"type": "json_object"},               },           )   def interactive_graph(kg: KnowledgeGraph) -> str:       net = Network(notebook=True, cdn_resources="remote")       node_map = {}       for index in range(len(kg.entities)):           node_map[kg.entities[index].name] = index           net.add_node(               index, label=kg.entities[index].name, title=kg.entities[index].description           )       for index in range(len(kg.relationships)):           relation = kg.relationships[index]           src_index = node_map[relation.source_entity]           target_index = node_map[relation.target_entity]           net.add_edge(src_index, target_index)       filename = "kg.html"       net.save_graph(filename)       return filename   def get_query_embedding(query: str):       zhipu_ai_client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))       response = zhipu_ai_client.embeddings.create(           model="embedding-2",           input=[query],       )       return response.data[0].embedding   def generate_result(query: str, entities, relationships):       zhipu_ai_client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))       entities_prompt = "\n".join(           map(lambda e: f'(Name: "{e.name}", Description: "{e.description}")', entities)       )       relationships_prompt = "\n".join(           map(lambda r: f'"{r.relationship_desc}"', relationships)       )       response = zhipu_ai_client.chat.completions.create(           model="glm-4-0520",           messages=[               {                   "role": "system",                   "content": "Please carefully think the user's "                   + "question and ONLY use the content below to generate answer:\n"                   + f"Entities: {entities_prompt}, Relationships: {relationships_prompt}",               },               {"role": "user", "content": query},           ],       )       return response.choices[0].message.content   def get_db_url():       return URL(           drivername="mysql+pymysql",           username="改成你的用户",           password="改成你的密码",           host="gateway01.us-west-2.prod.aws.tidbcloud.com",           port=4000,           database="chatdb",           query={"ssl_verify_cert": True, "ssl_verify_identity": True},       )   engine = create_engine(get_db_url(), pool_recycle=300)   Base = declarative_base()   class DatabaseEntity(Base):       id = Column(Integer, primary_key=True)       name = Column(String(512))       description = Column(Text)       description_vec = Column(VectorType(1024), comment="HNSW(distance=cosine)")       __tablename__ = "entities"   class DatabaseRelationship(Base):       id = Column(Integer, primary_key=True)       source_entity_id = Column(Integer, ForeignKey("entities.id"))       target_entity_id = Column(Integer, ForeignKey("entities.id"))       relationship_desc = Column(Text)       source_entity = relationship("DatabaseEntity", foreign_keys=[source_entity_id])       target_entity = relationship("DatabaseEntity", foreign_keys=[target_entity_id])       __tablename__ = "relationships"   def clean_knowledge_graph(kg: KnowledgeGraph) -> KnowledgeGraph:       entity_name_set = set(map(lambda e: e.name, kg.entities))       kg.relationships = list(           filter(               lambda r: r.source_entity in entity_name_set               and r.target_entity in entity_name_set,               kg.relationships,           )       )       return kg   def save_knowledge_graph(kg: KnowledgeGraph):       data_entities = list(           map(               lambda e: DatabaseEntity(                   name=e.name,                   description=e.description,                   description_vec=get_query_embedding(e.description),               ),               kg.entities,           )       )       with Session(engine) as session:           session.add_all(data_entities)           session.flush()           entity_id_map = dict(map(lambda e: (e.name, e.id), data_entities))           print(entity_id_map)           data_relationships = list(               map(                   lambda r: DatabaseRelationship(                       source_entity_id=entity_id_map[r.source_entity],                       target_entity_id=entity_id_map[r.target_entity],                       relationship_desc=r.relationship_desc,                   ),                   kg.relationships,               )           )           session.add_all(data_relationships)           session.commit()   def retrieve_entities_relationships(question_embedding) -> (List[DatabaseEntity], List[DatabaseRelationship]):       with Session(engine) as session:           entity = (               session.query(DatabaseEntity)               .order_by(                   DatabaseEntity.description_vec.cosine_distance(question_embedding)               )               .limit(1)               .first()           )           entities = {entity.id: entity}           relationships = (               session.query(DatabaseRelationship)               .options(                   joinedload(DatabaseRelationship.source_entity),                   joinedload(DatabaseRelationship.target_entity),               )               .filter(                   or_(                       DatabaseRelationship.source_entity == entity,                       DatabaseRelationship.target_entity == entity,                   )               )           )           for r in relationships:               entities.update(                   {                       r.source_entity.id: r.source_entity,                       r.target_entity.id: r.target_entity,                   }               )           return entities.values(), relationships   extractor = Extractor()   Base.metadata.create_all(engine)   zhipu_ai_client = dspy.OpenAI(model="glm-4-0520", api_base="https://open.bigmodel.cn/api/paas/v4/", api_key=os.getenv("ZHIPUAI_API_KEY"), model_type="chat", max_tokens=4096)   dspy.settings.configure(lm=zhipu_ai_client)   wiki = WikipediaLoader(query="TiDB").load()   pred = extractor(text=wiki[0].page_content)   knowledge_graph = clean_knowledge_graph(pred.knowledge)   interactive_graph(knowledge_graph)   save_knowledge_graph(knowledge_graph)   EOF   ls -l build-graph.py

$ cat > test-graph.py <<'EOF'   from zhipuai import ZhipuAI   import os   import click   from sqlalchemy import (       Column,       Integer,       String,       Text,       ForeignKey,       URL,       create_engine,       or_,   )   from typing import Mapping, Any, Optional, List   from sqlalchemy.orm import relationship, Session, declarative_base, joinedload   from tidb_vector.sqlalchemy import VectorType   def get_db_url():       return URL(           drivername="mysql+pymysql",           username="改成你的用户",           password="改成你的密码",           host="gateway01.us-west-2.prod.aws.tidbcloud.com",           port=4000,           database="chatdb",           query={"ssl_verify_cert": True, "ssl_verify_identity": True},       )   engine = create_engine(get_db_url(), pool_recycle=300)   Base = declarative_base()   class DatabaseEntity(Base):       id = Column(Integer, primary_key=True)       name = Column(String(512))       description = Column(Text)       description_vec = Column(VectorType(1024), comment="HNSW(distance=cosine)")       __tablename__ = "entities"   class DatabaseRelationship(Base):       id = Column(Integer, primary_key=True)       source_entity_id = Column(Integer, ForeignKey("entities.id"))       target_entity_id = Column(Integer, ForeignKey("entities.id"))       relationship_desc = Column(Text)       source_entity = relationship("DatabaseEntity", foreign_keys=[source_entity_id])       target_entity = relationship("DatabaseEntity", foreign_keys=[target_entity_id])       __tablename__ = "relationships"   def get_query_embedding(query: str):       zhipu_ai_client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))       response = zhipu_ai_client.embeddings.create(           model="embedding-2",           input=[query],       )       return response.data[0].embedding   def retrieve_entities_relationships(       question_embedding,   ) -> (List[DatabaseEntity], List[DatabaseRelationship]):       with Session(engine) as session:           entity = (               session.query(DatabaseEntity)               .order_by(                   DatabaseEntity.description_vec.cosine_distance(question_embedding)               )               .limit(1)               .first()           )           entities = {entity.id: entity}           relationships = (               session.query(DatabaseRelationship)               .options(                   joinedload(DatabaseRelationship.source_entity),                   joinedload(DatabaseRelationship.target_entity),               )               .filter(                   or_(                       DatabaseRelationship.source_entity == entity,                       DatabaseRelationship.target_entity == entity,                   )               )           )           for r in relationships:               entities.update(                   {                       r.source_entity.id: r.source_entity,                       r.target_entity.id: r.target_entity,                   }               )           return entities.values(), relationships   def generate_result(query: str, entities, relationships):       zhipu_ai_client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))       entities_prompt = "\n".join(           map(lambda e: f'(Name: "{e.name}", Description: "{e.description}")', entities)       )       relationships_prompt = "\n".join(           map(lambda r: f'"{r.relationship_desc}"', relationships)       )       response = zhipu_ai_client.chat.completions.create(           model="glm-4-0520",           messages=[               {                   "role": "system",                   "content": "Please carefully think the user's "                   + "question and ONLY use the content below to generate answer:\n"                   + f"Entities: {entities_prompt}, Relationships: {relationships_prompt}",               },               {"role": "user", "content": query},           ],       )       return response.choices[0].message.content   @click.command()   def start_chat():     while True:         question = click.prompt("Enter your question")         question_embedding = get_query_embedding(question)         entities, relationships = retrieve_entities_relationships(question_embedding)         result = generate_result(question, entities, relationships)         click.echo(result)   if __name__ == '__main__':     start_chat()   EOF   ls -l test-graph.py

$ python build-graph.py;python test-graph.py

   wiki = WikipediaLoader(query="TiDB").load()   pred = extractor(text=wiki[0].page_content)