pip install BeautifulSoup
pip install bs4

import os
DASHSCOPE_API_KEY=os.environ["DASHSCOPE_API_KEY"]
from langchain.embeddings import DashScopeEmbeddings
embeddings = DashScopeEmbeddings(model="text-embedding-v1", dashscope_api_key=DASHSCOPE_API_KEY)
from langchain_community.llms import Tongyi
llm_tongyi=Tongyi(temperature=1)

from langchain.schema import Document
from langchain_core.retrievers import BaseRetriever
class CustomRetriever(BaseRetriever):
    def __init__(self):
        super().__init__()
    def _get_relevant_documents(self, query: str) -> list[Document]:
        relevant_docs = []
        cur = conn.cursor()
        # Get the top 5 most similar documents using the KNN <=> operator
        cur.execute("SELECT content FROM docs ORDER BY embedding <=> '{}' LIMIT 5".format(embeddings.embed_query(query)))
        top3_docs = cur.fetchall()
        cur.close()
        
        for doc in top3_docs:
            relevant_docs.append(Document(page_content=doc[0]))
        return relevant_docs
custom_retriever = CustomRetriever()

from langchain.chains import RetrievalQA
# 创建 RetrievalQA
qa_chain = RetrievalQA.from_chain_type(
    llm=llm_tongyi,
    retriever=custom_retriever,
    verbose=True
)  
response = qa_chain.invoke("李士祎关连的股票信息?")
print(response)

from langchain.chains import GraphCypherQAChain
from langchain_community.graphs.age_graph import AGEGraph
conf = {
    "database": "xxx",
    "user":"xxx",
    "host":"pc-xxxx.rwlb.rds.aliyuncs.com",
    "port":"1921",
    "password":"xxx",
}
graph = AGEGraph(graph_name="stock_graph", conf=conf)
graph.refresh_schema()

graphchain = GraphCypherQAChain.from_llm(
    llm_tongyi, 
    graph=graph, 
    verbose=True, 
    top_k=5,
)
response = graphchain.invoke("李士祎关连的股票信息?")
print(response)

from langchain_core.prompts.prompt import PromptTemplate
from langchain.chains import LLMChain
graphchain = GraphCypherQAChain.from_llm(
    llm_tongyi, 
    graph=graph, 
    verbose=True, 
    top_k=5,
    return_direct=True, #此处直接返回查询结果，不需要llm来组织回答
)
template = """Task:Generate more detailed information about the raw answer.
The question is:
{question}
the raw answer from knowledge graph in format:
{answer}
please give more detailed answer about raw answer
"""
prompt = PromptTemplate(
    input_variables=["question", "answer"],
    template=template
)
llm_chain = LLMChain(llm=llm_tongyi, prompt=prompt)
def dynamic_query(question):
    # first round use graph cypher qa
    answer = graphchain.invoke(question)
    
    if answer['result'] == '':
        return "没有找到答案"
    # 格式化输入提示
    formatted_prompt = prompt.format(question=question,
                                    answer=answer['result'][0])
    
    answer = qa_chain.invoke(formatted_prompt)
    return answer
response = dynamic_query("李士祎关连的股票信息?")
print(response['result'])

import csv
import os
def remove_after_colon(input_string):
    # 查找冒号的位置
    colon_index = input_string.find(':')
    # 如果找到冒号，返回冒号之前的内容；否则返回原字符串
    if colon_index != -1:
        return input_string[:colon_index]
    return input_string
def convert_vertex_csv(file_path, graph, label):
    
    # create vlabel
    print("------------------------------------------------")
    print("-- Create vlabel")
    print("SELECT create_vlabel('{}','{}');".format(graph, label))
    with open(file_path, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter=',')
        header = next(reader)
        for row in reader:
            p = ""
            for h in header:
                h_name = h.lower()                
                if (h_name is None or h_name.startswith(":")):
                    continue;
                
                h_name = remove_after_colon(h_name)
                if p != "":
                    p += ","
                else:
                    p += "{"
                p += "{}:'{}'".format(h_name, row[header.index(h)].strip())
            if p != "":
               p += "}"
            print("SELECT * FROM cypher('{}', $$ CREATE (:{} {}) $$ ) as (n agtype);".format(graph, label, p))
                
def convert_edge_csv(file_path, graph, from_type, to_type):
    file_name = os.path.splitext(os.path.basename(file_path))[0].lower()
    with open(file_path, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter=',')
        header = next(reader)
        for row in reader:
            p = ""
            REL_TYPE="RELTYPE"
            START_ID=""
            END_ID=""
            for h in header:                
                h_name = h.lower()
                if (h_name is None):
                    continue;
                elif h_name.startswith(":start_id"):
                    START_ID=row[header.index(h)].strip()
                    continue
                elif h_name.startswith(":end_id"):
                    END_ID=row[header.index(h)].strip()
                    continue
                elif h_name.startswith(":type"):
                    REL_TYPE=row[header.index(h)].strip()
                    continue
                if p != "":
                    p += ","
                else:
                    p += "{"
                p += "{}:'{}'".format(h.lower(), row[header.index(h)].strip())
            
            if p != "":
               p += "}"
            
            print("SELECT * FROM cypher('{0}', $$ MATCH (a:{1}), (b:{2}) WHERE a.{1}_id = '{3}' AND "
                  "b.{2}_id = '{4}' CREATE (a)-[e:{6} {5} ]->(b) RETURN e$$) as (e agtype);".format(graph, from_type, to_type, START_ID, END_ID, p, REL_TYPE))
# convert vertex
convert_vertex_csv('data/import/executive.csv','stock_graph', 'person')
# convert edge
convert_edge_csv('data/import/stock_concept.csv','stock_graph', 'stock', 'concept')

from langchain.text_splitter import RecursiveCharacterTextSplitter
from bs4 import BeautifulSoup
import os
from langchain_community.embeddings import DashScopeEmbeddings
import psycopg
connection = "dbname=xxx user=xxx host=xxx port=1921 password=xxxx"  
# DASHSCOPE_API_KEY从环境变量中读取
DASHSCOPE_API_KEY=os.environ["DASHSCOPE_API_KEY"]
embeddings = DashScopeEmbeddings(model="text-embedding-v1", dashscope_api_key=DASHSCOPE_API_KEY)
conn = psycopg.connect(connection)
cur = conn.cursor()
table_create_command = """
CREATE TABLE IF NOT EXISTS docs (
            id bigserial primary key, 
            title text,
            content text,
            tokens integer,
            embedding vector(1536)
            );
            """
cur.execute(table_create_command)
cur.close()
conn.commit()
def split_doc(file):
    html_content=None
    with open(file, 'r', encoding='GBK') as f:
        html_content = f.read()
    soup = BeautifulSoup(html_content, 'html.parser')
    text_content = soup.get_text(strip=True)  # 去掉多余的空白字符
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = text_splitter.split_text(text_content)
    for chunk in chunks:
        cur = conn.cursor()
        sql = "INSERT INTO docs (title, content, tokens, embedding) VALUES ('{}', '{}', {}, '{}')".format(os.path.basename(file), chunk.replace("'", "''"), len(chunk), embeddings.embed_query(chunk))
        cur.execute(sql)
        cur.close()
        conn.commit()
# 打印每个切分后的文档
def get_all_files(directory):
    file_paths = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_paths.append(os.path.join(root, file))
    return file_paths
directory_path = '/path/stockpage/'
all_files = get_all_files(directory_path)
for file in all_files:
    print(file)
    if file.endswith('.html'):
        split_doc(file)